EGraph

A generic slotted e-graph engine in Swift — represent many equivalent forms of a term at once, rewrite them to a fixpoint, and extract the best one.

Start with a language built on the engine — the Algebra example is an optimizer and a linear-equation solver:

import Algebra

let algebra = AlgebraEngine()
algebra.optimize("x * 2")          // "x << 1"   — strength reduction, chosen by a cost model
algebra.simplify("(x + 0) * 1")    // "x"        — identities + constant folding
algebra.solve("2 * x + 3 = 7")     // "x = 2"    — a linear-equation solver

Here is how it's built — and the shape of any language you'd write on the engine. Declare operators with @Language, give them compile-checked rewrite rules and a cost, then saturate and extract the best form:

import EGraphDSL

// `@Language` derives the operator/arity table from the `@arity` markers — the whole language definition
// (and adds the Hashable/Sendable/BitwiseCopyable conformances the engine needs):
@Language enum Algebra {
    @arity(0) case number(Int), x
    @arity(2) case add, subtract, multiply, shiftLeft, equals
}

// Rewrite rules, checked against Algebra's operators at compile time:
let rules = RuleSet<Algebra> {
    Rewrite("strength-reduction",
            from: pattern(.multiply, metavariable(0), leaf(.number(2))),   // x * 2
            to:   pattern(.shiftLeft, metavariable(0), leaf(.number(1))))  // x << 1
    Rewrite("add-identity",
            from: pattern(.add, metavariable(0), leaf(.number(0))),        // x + 0
            to:   metavariable(0))                                          // x
}

// A cost model — a shift is cheaper than a multiply, so extraction prefers `x << 1` over `x * 2`:
struct Cost: CostFn {
    typealias L = Algebra
    func cost(_ op: Algebra, children: [Int]) -> Int { (op == .multiply ? 4 : 1) + children.reduce(0, +) }
}

// Saturate to a fixpoint, then extract the cheapest form. (`parse` is a small parser you add per language —
// the Algebra example ships one as an EGraph extension; `simplify`/`solve` are more rules over the same setup.)
var g = EGraph<Algebra, NoAnalysis<Algebra>, NoSlots>()
let expr = g.parse("x * 2")
_ = g.saturate(rules)
g.extract(expr, using: Cost())        // ⇒  x << 1   (shiftLeft(x, number(1)))

See Build your own language for the full walkthrough.

Because the engine is slotted, the Lambda example gets binders for free — α-equivalence and capture-avoiding substitution come from the data structure, so λx. x and λy. y are the same e-class:

import Lambda

let lambda = LambdaEngine()
lambda.equal("λx. x", "λy. y")             // true      — same e-class, no renaming pass
lambda.normalize("(λx. x) (λy. y)").term   // "λx. x"   — β-reduction

And because instructions have a cost, the same cost-based extraction is a peephole superoptimizer. The ARM64 example takes a program of actual Apple-Silicon assembly and emits the optimal instruction sequence — strength reduction, instruction selection (fusing madd / shifted-register operands), constant folding, dead-code elimination, and CSE all at once. A bloated 14-instruction naive codegen for 8*a + b collapses to a single instruction:

import ARM64

let arm = ARM64Optimizer()
try arm.compile("""
mul  x2, x0, #8      // ┐ four multiplies, redundant recomputation,
add  x3, x2, x1      // │ dead instructions (x4–x8), and identity ops
mul  x4, x0, #8      // │ (×1, +0, −0, x^x) …
add  x5, x4, x1
mul  x6, x0, #0
mul  x7, x1, #1
mul  x8, x0, x1
add  x9, x3, x6
mul  x10, x9, #1
add  x11, x10, #0
sub  x12, x11, #0
eor  x13, x12, x12
add  x14, x11, x13
mov  x15, x14
""")
// ⇒  "add x9, x1, x0, lsl #3"      // …all of it collapses to one shifted-register add (8a + b)

Swap the cost model and the same machinery is a boolean minimizer — here "cost" is gate count, so the Logic example extracts the fewest-gate equivalent:

import Logic

let logic = Logic()
logic.minimize("a | (a & b)")   // "a"   — absorption
logic.minimize("a & !a")        // "0"   — complement
logic.minimize("!(!a)")         // "a"   — double negation

These are small example languages built on the engine — you write your own the same way (see Build your own language).

• EGraph — the engine (hashcons + union-find + congruence closure, equality saturation, cost-based extraction, slotted binders). No dependencies.
• EGraphDSL — the @Language macro + a typesafe RuleSet builder (depends on swift-syntax).
• Lambda, Algebra, ARM64, Logic — each example language (under Examples/) is its own target, so importing one does not pull the others.

What is an e-graph good for?

Equality saturation shines whenever you have many equivalent ways to write something and want the best one, or you want to prove two things equal without committing to a rewrite order. Beyond these examples, the same engine could underpin:

• Compiler optimization — a peephole/IR optimizer that finds the cheapest equivalent code: strength reduction, reassociation, constant folding, redundant-load elimination.
• Floating-point accuracy — rewrite a numeric expression into an equivalent one with lower rounding error.
• Tensor / array graph rewriting — fuse and re-layout operators ((AB)C = A(BC), Aᵀᵀ = A), extracting by estimated cost.
• Database query optimization — predicate pushdown, join reordering; extract the plan with the lowest estimated cost.
• Symbolic math — simplification and, with a d/dx binder, symbolic differentiation — a natural fit for the slotted engine.
• Regular-expression / boolean-algebra simplification — normalize to the smallest equivalent form.
• Equality-modulo-rules kernels — deciding term equality under user-supplied rewrite rules, where the slotted representation handles binders.

The shipped examples are deliberately small so the whole pipeline is legible: source string → parse → e-graph → saturate → extract → print.

Build your own language

A language is: an operator set, an optional analysis, some rewrite rules, and (usually) a parser + printer. Here is the shape, using the Algebra example.

1. Declare the operators. @Language derives the Language conformance (the arity(of:) table) from the @arity markers — no hand-written, drift-prone switch — and adds the operator conformances the engine needs (Hashable, Sendable, BitwiseCopyable), so you don't repeat them:

import EGraphDSL

@Language
public enum Algebra {
    @arity(0) case number(Int)   // an integer literal
    @arity(0) case x             // the unknown
    @arity(2) case add, subtract, multiply, divide, shiftLeft, equals
}

2. (Optional) an Analysis — a semilattice fact per e-class, computed bottom-up. Here, constant folding:

struct ConstantFolding: Analysis {
    typealias L = Algebra; typealias Data = Int?          // the class's known value, or nil
    func make(op: L.Op, children: [Data]) -> Data { /* fold number/add/multiply/… */ }
    mutating func join(_ into: inout Data, _ other: Data) -> DidChange { /* nil → known */ }
}

3. Rewrite rules — the typesafe RuleSet builder. Operator names are checked by the compiler (they must be real Algebra cases) and each constructor enforces its arity by its signature:

let optimizeRules = RuleSet {
    Rewrite("multiply-by-2", from: multiply(a, number(2)), to: shiftLeft(a, number(1)))   // strength reduction
    Rewrite("add-zero", from: add(a, number(0)), to: a)                // identity
}

where a = metavariable(0) and multiply/add/number are one-line typed constructors for your operators. Rules live wherever you like (an init, a property) — they are a RuleSet<Algebra> (usable as a collection of rules).

4. Saturate + extract. Cost-based extraction pulls out the cheapest equivalent term:

var g = EGraph<Algebra, ConstantFolding, NoSlots>(analysis: ConstantFolding())
let root = g.add(/* parsed program */)
_ = g.saturate(optimizeRules)
let best = g.extract(root, using: myCostModel)   // e.g. shift < add < multiply  ⟹  x * 2 becomes x << 1

Programs themselves are plain strings parsed by the language's own recursive-descent parser (see the Algebra target), so users write "2 * x + 3 = 7", not Swift.

Rules: compile-time typed, or runtime strings

The typed RuleSet builder is the recommended path: operator names and arities are checked by the Swift compiler, so a mistake in a rule is a compile error. String rules are also supported (see AlgebraEngine.parseRule) and validated at runtime — handy when rules come from data or config rather than source.

The examples (what each one teaches)

• Lambda — untyped λ-calculus, slotted. α-equivalence is automatic and β is capture-avoiding; reduction is maxRounds-bounded and returns a StopReason (the engine promises bounded, non-hanging work — not totality). A neat consequence: the e-graph saturates on Ω = (λx. x x)(λx. x x), since the reduct is α-equal to the redex (same class).
• Algebra — expression optimizer + linear-equation solver, non-slotted. Shows a ConstantFolding Analysis (reified into number nodes), cost-based extraction (strength reduction), and that permutative rules (commutativity/associativity) saturate instead of looping.
• ARM64 — an Apple-Silicon assembly superoptimizer, non-slotted. Assembly in, optimal assembly out: strength reduction (mul ×5 → add x, x, x, lsl #2, matching real GCC/LLVM), instruction selection (fusing madd/msub and shifted-register add/sub/and/orr/eor), constant folding, dead-code elimination, and CSE (a shared subexpression is emitted once). Cost = approximate instruction latency.
• Logic — a boolean-logic minimizer, non-slotted. "Cost" is gate count: idempotence, absorption, complement, and constant folding; extraction pulls out the fewest-gate equivalent.

Programs are plain strings parsed by each language's own parser — the honest pipeline: source → parse → e-graph → saturate → extract → print.

Build & test

swift build
swift test

EGraph has no dependencies; EGraphDSL (the @Language/RuleSet layer) uses swift-syntax for the macro. Each example and the DSL have their own test target.

References

• Schneider, Rossel, Shaikhha, Goens, Koehler, Steuwer. Slotted E-Graphs (PLDI 2025).
• Willsey, Nandi, Wang, Flatt, Tatlock, Panchekha. egg: Fast and Extensible Equality Saturation (POPL 2021).