Skip to content

Implement egg-style deferred rebuilding algorithm #158

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 2 commits into from
Dec 22, 2025
Merged

Implement egg-style deferred rebuilding algorithm #158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
08dd04e
Initial plan
Copilot Dec 22, 2025
d2df311
Implement egg-style deferred rebuilding algorithm
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
56 changes: 34 additions & 22 deletions egg/apyds_egg/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ def __init__(self) -> None:
self.unionfind: UnionFind[EClassId] = UnionFind()
self.classes: dict[EClassId, set[ENode]] = {}
self.parents: dict[EClassId, set[tuple[ENode, EClassId]]] = defaultdict(set)
self.worklist: set[EClassId] = set()

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

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

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

self._repair(r)
self.worklist.add(r)

return r

def _repair(self, eclass: EClassId) -> None:
"""Immediately restore congruence by re-canonicalizing parents and merging congruent ones.
def rebuild(self) -> None:
"""Restore congruence by processing the worklist.

This method implements the traditional E-Graph repair algorithm:
This method implements the egg-style deferred rebuilding:
- Process all E-classes in the worklist
- Re-canonicalize parents and merge congruent ones
- Continue until worklist is empty
"""
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 repair(self, eclass: EClassId) -> None:
"""Restore congruence for a single E-class.

This method implements the egg-style repair algorithm:
- Re-canonicalize all parent nodes
- Merge congruent parents recursively
- Continue until no more changes occur
- Merge congruent parents (which may add more work to worklist)
- Update hashcons and parent tracking
"""
changed = True
while changed:
changed = False
new_parents: dict[ENode, EClassId] = {}
new_parents: dict[ENode, EClassId] = {}

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

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

if canon in new_parents:
self.merge(peclass, new_parents[canon])
changed = True
else:
new_parents[canon] = peclass
self.hashcons[canon] = peclass
if canon in new_parents:
self.merge(peclass, new_parents[canon])
Comment on lines +209 to +210

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 repair() method is calling self.merge() at line 210, which adds items to self.worklist. However, since repair() is called from within rebuild() which is already processing the worklist, this could lead to correctness issues if the worklist handling isn't carefully managed. While the current implementation appears correct (as the worklist is cleared and rebuilt in each iteration), the recursive nature of calling merge() from within repair() should be clearly documented. Consider adding a comment explaining that merge operations during repair are safe because the worklist is rebuilt in each iteration of rebuild().

Copilot uses AI. Check for mistakes.
else:
new_parents[canon] = peclass
self.hashcons[canon] = peclass

self.parents[eclass] = {(p, c) for p, c in new_parents.items()}
self.parents[eclass] = {(p, c) for p, c in new_parents.items()}
3 changes: 3 additions & 0 deletions egg/tests/test_egraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ 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 @@ -157,6 +158,7 @@ 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 Down Expand Up @@ -195,6 +197,7 @@ 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