sml-egraph

An e-graph (equivalence graph) in pure Standard ML: a union-find-backed data structure over expression terms that compactly represents many equivalent programs at once, plus equality saturation -- applying rewrite rules to all of them simultaneously until fixpoint -- and cost-based extraction of the best representative. This is the core data structure behind term-rewriting/optimization engines such as egg; terminology follows "egg: Fast and Extensible Equality Saturation" (Willsey, Nandi, Wang, Flatt, Tatlock, Panchekha; POPL 2021).

No dependencies, no FFI, no threads, no clock, no randomness: the same sequence of operations always produces the same result under MLton and Poly/ML. The egraph type is an internally-mutable handle (union-find + hashcons table implemented with plain, deterministic association lists -- no hashing, so no cross-compiler Word/hash-order divergence); functions that "update" an egraph mutate it in place and return it so callers can still thread it through let val (g, id) = addTerm g t, but -- unlike this fleet's usual persistent-tree libraries (e.g. sml-fenwick) -- an earlier binding is not an independent snapshot once later operations mutate the same underlying graph.

• Core: e-nodes (ENode of string * id list, an operator applied to child e-class ids), e-classes (implicit: a set of e-nodes sharing a union-find root id), a hashcons table for structural dedup, and find/merge/rebuild maintaining the congruence invariant -- if two e-nodes have the same operator and pairwise-equivalent children, their e-classes are merged.
• Equality saturation: match rewrite-rule patterns (PVar/PNode) against every e-class, instantiate the right-hand side for every match, merge, and rebuild -- repeated to fixpoint or an iteration cap.
• Extraction: pull out the minimum-cost representative term of an e-class under a caller-supplied per-operator cost function, via a Bellman-Ford-style relaxation to fixpoint.

rebuild restores congruence with a full-table fixpoint pass rather than egg's incremental worklist-over-parents algorithm -- a deliberate scope cut for implementation simplicity; it produces the identical congruence-closed result, just less asymptotically efficient (fine at the scale this library targets). See egraph.sig for the full documented API.

API

signature EGRAPH =
sig
  type id = int
  datatype term  = Node of string * term list
  datatype enode = ENode of string * id list
  datatype pat   = PVar of string | PNode of string * pat list
  type rule = { name : string, lhs : pat, rhs : pat }
  type egraph

  val empty : unit -> egraph

  val find         : egraph -> id -> id
  val equivalent   : egraph -> id -> id -> bool
  val canonicalize : egraph -> enode -> enode
  val addEnode     : egraph -> enode -> egraph * id
  val addTerm      : egraph -> term -> egraph * id
  val merge        : egraph -> id -> id -> egraph
  val rebuild      : egraph -> egraph

  val classes    : egraph -> id list
  val classNodes : egraph -> id -> enode list
  val numClasses : egraph -> int
  val numNodes   : egraph -> int

  val applyRules : rule list -> egraph -> egraph * bool
  val saturate   : rule list -> int -> egraph -> egraph * int

  exception UnboundPatternVar of string
  exception NoExtraction
  val extract : (string -> int) -> egraph -> id -> term
end

Example

val g0 = Egraph.empty ()
val (g1, idA)      = Egraph.addTerm g0 (Egraph.Node ("a", []))
val (g2, idAPlus0) = Egraph.addTerm g1 (Egraph.Node ("+", [Egraph.Node ("a", []), Egraph.Node ("0", [])]))

val identityRule =
  { name = "add-zero",
    lhs  = Egraph.PNode ("+", [Egraph.PVar "x", Egraph.PNode ("0", [])]),
    rhs  = Egraph.PVar "x" }

val (g3, changed) = Egraph.applyRules [identityRule] g2
val true = Egraph.equivalent g3 idA idAPlus0

val term = Egraph.extract (fn _ => 1) g3 idAPlus0   (* => Node ("a", []) *)

Running examples/demo.sml with make example prints:

1. Hashcons dedup: adding (+ a 0) twice
   same id both times: yes
   e-classes so far: 3 (a, 0, (+ a 0))

2. Congruence: f(a,b) and f(a,c) merge once b = c
   f(a,b) = f(a,c) before merging b,c? no
   f(a,b) = f(a,c) after merging b,c?  yes

3. Rewrite rules: (+ ?x 0) -> ?x and (+ ?x ?y) -> (+ ?y ?x)
   a = (+ a 0) after one pass? yes (changed: yes)

4. Saturation: (a * 2) / 2  ->  a
   starting term: (/ (* a 2) 2)
   passes to fixpoint: 1
   equivalent to 'a' after saturation? yes

5. Extraction: minimum-cost representative of a class
   extract (a*2)/2's class, unit cost => a
   extract (+ a 0)'s class, unit cost => a
   same class {a,z}, cost prefers 'a' expensive => z
   same class {a,z}, cost prefers 'z' expensive => a

Build & test

Requires MLton and/or Poly/ML.

make test        # build + run the suite under MLton
make test-poly   # run the suite under Poly/ML
make all-tests   # both + byte-identical diff of their stdout
make example     # build + run the demo
make clean

Installing with smlpkg

smlpkg add github.com/sjqtentacles/sml-egraph
smlpkg sync

Reference lib/github.com/sjqtentacles/sml-egraph/egraph.mlb from your own .mlb (MLton / MLKit), or feed sources.mlb to tools/polybuild (Poly/ML).

Layout

sml.pkg                                       smlpkg manifest
Makefile                                      MLton + Poly/ML targets
.github/workflows/ci.yml                      CI: MLton + Poly/ML
lib/github.com/sjqtentacles/sml-egraph/
  egraph.sig     EGRAPH signature
  egraph.sml     union-find + hashcons + congruence + saturation + extraction
  sources.mlb    ordered source list
  egraph.mlb     public basis
examples/
  demo.sml       hashcons/congruence/rewrite/saturation/extraction walkthrough
test/
  harness.sml    shared assertion harness
  test.sml       44 checks against a small fixed term language
  entry.sml / main.sml
tools/polybuild  Poly/ML build wrapper

Tests

44 deterministic checks: hashcons structural dedup on addTerm/addEnode; union-find find/merge/equivalent including idempotent re-merging and transitivity; congruence propagation through one and two levels of nested terms (f(a,b)/f(a,c) merging when b/c merge, cascading to g(f(a,b),x)/g(f(a,c),x)); rewrite-rule application (additive identity, commutativity, multiplicative identity, and a no-match case); equality saturation reaching a fixpoint before an iteration cap on (a * 2) / 2 -> a; minimum-cost extraction under both a unit cost function and custom weighted cost functions that must change which representative is picked; and error paths (NoExtraction on an unknown class, UnboundPatternVar on a rewrite rule whose right-hand side references a variable its left-hand side never binds). Run make all-tests to verify identical output under both compilers.

License

MIT. See LICENSE.