Skip to content

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

Merged
hzhangxyz merged 3 commits into
mainfrom
copilot/use-original-egraph-implementation
Dec 22, 2025
Merged

Replace egg-style deferred rebuilding with traditional immediate E-Graph repair#154
hzhangxyz merged 3 commits into
mainfrom
copilot/use-original-egraph-implementation

Conversation

Copilot AI commented Dec 22, 2025

Copy link
Copy Markdown
Contributor

The current implementation uses egg-style deferred rebuilding with a worklist, allowing temporary invariant violations. This PR replaces it with the traditional E-Graph algorithm that maintains congruence immediately.

Changes

  • Removed worklist-based deferred scheduling: Deleted worklist attribute from EGraph.__init__()
  • Modified merge() to repair immediately: Now calls _repair(r) directly after merging instead of adding to worklist
  • Updated _repair() with iterative convergence: Added changed flag and while loop to handle cascading upward merges until congruence is fully restored
  • Removed rebuild() method entirely: Method is no longer needed since congruence is maintained immediately during merge()
  • Updated all tests: Removed all rebuild() calls from tests and renamed tests that referenced rebuilding

Behavior difference

eg = EGraph()
a = eg.add(apyds.Term("a"))
b = eg.add(apyds.Term("b"))
ax = eg.add(apyds.Term("(+ a x)"))
bx = eg.add(apyds.Term("(+ b x)"))

# Before: needed rebuild() to restore congruence
eg.merge(a, b)
eg.rebuild()  # Required!
assert eg.are_equal(ax, bx)

# After: congruence maintained immediately, no rebuild() method
eg.merge(a, b)
assert eg.are_equal(ax, bx)  # Already true

Invariants are never violated. All tests pass with the rebuild() method completely removed.

Original prompt

This section details on the original issue you should resolve

<issue_title>使用原始的egraph实现</issue_title>
<issue_description>目前的egg实现参考的是:

# egraph_demo_typed_py313.py
# Minimal, typed E-Graph demo with deferred rebuilding (egg-style)
# Python 3.13+

from __future__ import annotations

from dataclasses import dataclass
from typing import NewType, Callable
from collections import defaultdict


# ---------- Strongly-typed IDs ----------

EClassId = NewType("EClassId", int)


# ---------- Union-Find ----------

class UnionFind:
    parent: dict[EClassId, EClassId]

    def __init__(self) -> None:
        self.parent = {}

    def find(self, x: EClassId) -> EClassId:
        if x not in self.parent:
            self.parent[x] = x
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a: EClassId, b: EClassId) -> EClassId:
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            self.parent[rb] = ra
        return ra


# ---------- ENode ----------

@dataclass(frozen=True)
class ENode:
    op: str
    children: tuple[EClassId, ...]

    def canonicalize(
        self,
        find: Callable[[EClassId], EClassId],
    ) -> ENode:
        return ENode(
            self.op,
            tuple(find(c) for c in self.children),
        )


# ---------- EGraph ----------

class EGraph:
    uf: UnionFind
    next_id: int

    classes: dict[EClassId, set[ENode]]
    parents: dict[EClassId, set[tuple[ENode, EClassId]]]
    hashcons: dict[ENode, EClassId]

    worklist: set[EClassId]

    def __init__(self) -> None:
        self.uf = UnionFind()
        self.next_id = 0

        self.classes = {}
        self.parents = defaultdict(set)
        self.hashcons = {}

        self.worklist = set()

    # ----- basic ops -----

    def _fresh_id(self) -> EClassId:
        eid = EClassId(self.next_id)
        self.next_id += 1
        return eid

    def find(self, eclass: EClassId) -> EClassId:
        return self.uf.find(eclass)

    def add(self, enode: ENode) -> EClassId:
        enode = enode.canonicalize(self.find)

        if enode in self.hashcons:
            return self.find(self.hashcons[enode])

        eid = self._fresh_id()

        self.uf.parent[eid] = eid
        self.classes[eid] = {enode}
        self.hashcons[enode] = eid

        for c in enode.children:
            self.parents[c].add((enode, eid))

        return eid

    def merge(self, a: EClassId, b: EClassId) -> EClassId:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return ra

        r = self.uf.union(ra, rb)

        # merge e-nodes
        self.classes[r] |= self.classes[rb]
        del self.classes[rb]

        # merge parent info
        self.parents[r] |= self.parents[rb]
        del self.parents[rb]

        # defer congruence restoration
        self.worklist.add(r)

        return r

    # ----- rebuilding (egg-style) -----

    def rebuild(self) -> None:
        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:
        new_parents: dict[ENode, EClassId] = {}

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

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

            if canon in new_parents:
                # congruent parents => upward merge
                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()}

    # ----- debug helpers -----

    def dump(self) -> None:
        print("E-Graph:")
        for eid, nodes in self.classes.items():
            print(f"  class {eid}:")
            for n in nodes:
                print(f"    {n}")
        print()


# ---------- Demo ----------

if __name__ == "__main__":
    eg = EGraph()

    # constants
    x = eg.add(ENode("x", ()))
    a = eg.add(ENode("a", ()))
    b = eg.add(ENode("b", ()))

    ax = eg.add(ENode("+", (a, x)))
    bx = eg.add(ENode("+", (b, x)))

    print("Before merge:")
    print("  (+ a x) eclass =", eg.find(ax))
    print("  (+ b x) eclass =", eg.find(bx))
    eg.dump()

    # assert a ≡ b
    eg.merge(a, b)
    eg.rebuild()

    print("After merge + rebuild:")
    print("  (+ a x) eclass =", eg.find(ax))
    print("  (+ b x) eclass =", eg.find(bx))
    eg.dump()

这个egraph实现使用了一些技巧,即 worklist / rebuild / deferred scheduling,

我现在希望使用原始的egraph:

  • 不允许 invariants 暂时被破坏
  • 没有 worklist / rebuild / deferred scheduling

我拿gpt写了一个demo:

# egraph_traditional_py313.py
# Traditional E-Graph with immediate upward merging
# Python 3....

</details>



<!-- START COPILOT CODING AGENT SUFFIX -->

- Fixes USTC-KnowledgeComputingLab/ds#153

<!-- START COPILOT CODING AGENT TIPS -->
---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our [2 minute survey](https://gh.io/copilot-coding-agent-survey).

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
Copilot AI changed the title [WIP] Implement original E-Graph from provided specification Replace egg-style deferred rebuilding with traditional immediate E-Graph repair Dec 22, 2025
Copilot AI requested a review from hzhangxyz December 22, 2025 05:12

@hzhangxyz hzhangxyz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

rebuild函数可以直接删了

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>

Copilot AI commented Dec 22, 2025

Copy link
Copy Markdown
Contributor Author

rebuild函数可以直接删了

Done. Removed the rebuild() method and updated all tests accordingly. (ef893ab)

@hzhangxyz
hzhangxyz marked this pull request as ready for review December 22, 2025 05:29
Copilot AI review requested due to automatic review settings December 22, 2025 05:29
@hzhangxyz
hzhangxyz merged commit fa64c02 into main Dec 22, 2025
83 checks passed

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

This PR replaces the egg-style deferred rebuilding mechanism with a traditional E-Graph algorithm that maintains congruence invariants immediately during merge operations, eliminating the need for explicit rebuild() calls.

Key Changes:

  • Removed the worklist-based deferred scheduling mechanism from EGraph
  • Modified merge() to call _repair() immediately after merging e-classes
  • Redesigned _repair() with an iterative convergence loop to handle cascading upward merges until full congruence is restored

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
egg/apyds_egg/init.py Removed worklist attribute, eliminated rebuild() method, modified merge() to call _repair() immediately, and updated _repair() to use iterative convergence with a while changed loop for handling cascading merges
egg/tests/test_egraph.py Removed all rebuild() calls after merge operations, renamed test_egraph_rebuild_empty_worklist to test_egraph_immediate_congruence, and renamed test_egraph_are_equal_after_rebuild to test_egraph_are_equal_after_merge to reflect immediate congruence maintenance

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread egg/tests/test_egraph.py


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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants