Skip to content

Replace egg-style deferred rebuilding with traditional immediate E-Graph repair #154

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hzhangxyz merged 3 commits into from
Dec 22, 2025
Merged

Replace egg-style deferred rebuilding with traditional immediate E-Graph repair #154

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3abe379
Initial plan
Copilot Dec 22, 2025
6ba86c1
Convert to traditional E-Graph with immediate repair
Copilot Dec 22, 2025
ef893ab
Remove rebuild() method as it's no longer needed
Copilot Dec 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 32 additions & 26 deletions egg/apyds_egg/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ def __init__(self) -> None:
self.classes: dict[EClassId, set[ENode]] = {}
self.parents: dict[EClassId, set[tuple[ENode, EClassId]]] = defaultdict(set)
self.hashcons: dict[ENode, EClassId] = {}
self.worklist: set[EClassId] = set()

def _fresh_id(self) -> EClassId:
"""Generate a fresh E-class ID."""
Expand Down Expand Up @@ -141,7 +140,7 @@ def _add_enode(self, enode: ENode) -> EClassId:
return eid

def merge(self, a: EClassId, b: EClassId) -> EClassId:
"""Merge two E-classes and schedule rebuilding.
"""Merge two E-classes and immediately restore congruence.

Args:
a: The first E-class ID to merge.
Expand All @@ -162,19 +161,11 @@ def merge(self, a: EClassId, b: EClassId) -> EClassId:
self.parents[r] |= self.parents[rb]
del self.parents[rb]

self.worklist.add(r)
# Immediately restore invariants
self._repair(r)

return r

def rebuild(self) -> None:
"""Restore congruence by processing the worklist."""
while self.worklist:
todo: set[EClassId] = {self.find(e) for e in self.worklist}
self.worklist.clear()

for eclass in todo:
self._repair(eclass)

def are_equal(self, a: EClassId, b: EClassId) -> bool:
"""Check if two E-class IDs are equivalent.

Expand All @@ -188,19 +179,34 @@ def are_equal(self, a: EClassId, b: EClassId) -> bool:
return self.find(a) == self.find(b)

def _repair(self, eclass: EClassId) -> None:
"""Repair congruence for an E-class by updating parent nodes."""
new_parents: dict[ENode, EClassId] = {}

for pnode, peclass in list(self.parents[eclass]):
self.hashcons.pop(pnode, None)
"""Immediately restore congruence by re-canonicalizing parents and merging congruent ones.

canon = pnode.canonicalize(self.find)
peclass = self.find(peclass)
This method implements the traditional E-Graph repair algorithm:
- Re-canonicalize all parent nodes
- Merge congruent parents recursively
- Continue until no more changes occur

if canon in new_parents:
self.merge(peclass, new_parents[canon])
else:
new_parents[canon] = peclass
self.hashcons[canon] = peclass

self.parents[eclass] = {(p, c) for p, c in new_parents.items()}
Args:
eclass: The E-class ID to repair.
"""
changed = True
while changed:
changed = False
new_parents: dict[ENode, EClassId] = {}

for pnode, peclass in list(self.parents[eclass]):
# Remove old hashcons entry
self.hashcons.pop(pnode, None)

canon = pnode.canonicalize(self.find)
peclass = self.find(peclass)

if canon in new_parents:
# Upward merge required
self.merge(peclass, new_parents[canon])
changed = True
else:
new_parents[canon] = peclass
self.hashcons[canon] = peclass

self.parents[eclass] = {(p, c) for p, c in new_parents.items()}
13 changes: 3 additions & 10 deletions egg/tests/test_egraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,6 @@ def test_egraph_congruence():
assert eg.find(ax) != eg.find(bx)

eg.merge(a, b)
eg.rebuild()

assert eg.find(ax) == eg.find(bx)

Expand All @@ -158,7 +157,6 @@ def test_egraph_congruence_nested():
bcc = eg.add(apyds.Term("(g (f b c) c)"))

eg.merge(a, b)
eg.rebuild()

assert eg.find(ac) == eg.find(bc)
assert eg.find(acc) == eg.find(bcc)
Expand All @@ -177,12 +175,10 @@ def test_egraph_multiple_merges():
assert eg.find(a) == eg.find(c)


def test_egraph_rebuild_empty_worklist():
def test_egraph_immediate_congruence():

Copilot AI Dec 22, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new test name test_egraph_immediate_congruence doesn't accurately describe what this test does. The test simply adds a term and verifies that find returns the same ID - it doesn't test congruence at all (congruence refers to structurally identical terms being merged automatically).

Consider renaming this test to something more descriptive like test_egraph_find_returns_self or test_egraph_single_term to better reflect its purpose.

Suggested change
def test_egraph_immediate_congruence():
def test_egraph_find_returns_self():

Copilot uses AI. Check for mistakes.
eg = EGraph()
a = eg.add(apyds.Term("a"))

eg.rebuild()

assert eg.find(a) == a


Expand All @@ -199,7 +195,6 @@ def test_egraph_complex_example():
assert eg.find(ax) != eg.find(bx)

eg.merge(a, b)
eg.rebuild()

assert eg.find(ax) == eg.find(bx)

Expand Down Expand Up @@ -251,7 +246,6 @@ def test_egraph_associativity_example():
assert eg.find(xy_z) != eg.find(x_yz)

eg.merge(xy_z, x_yz)
eg.rebuild()

assert eg.find(xy_z) == eg.find(x_yz)

Expand Down Expand Up @@ -302,7 +296,7 @@ def test_egraph_are_equal():
assert eg.are_equal(a, b)


def test_egraph_are_equal_after_rebuild():
def test_egraph_are_equal_after_merge():
eg = EGraph()

x = eg.add(apyds.Term("x"))
Expand All @@ -317,7 +311,6 @@ def test_egraph_are_equal_after_rebuild():

# Merge a and b
eg.merge(a, b)
eg.rebuild()

# After rebuild, ax and bx should be equal due to congruence
# After merge, ax and bx should be equal due to congruence
assert eg.are_equal(ax, bx)