Implement egg-style deferred rebuilding algorithm

[Copilot]

Replaces immediate upward merging with egg-style deferred rebuilding for better performance and explicit control over congruence restoration.

Changes

• Added worklist: Tracks E-classes requiring congruence restoration instead of immediate repair
• Modified merge(): Defers repair by adding to worklist rather than calling repair immediately
• Added rebuild(): New public method processes worklist until congruence is restored
• Refactored _repair() to repair(): Removed internal while changed: loop; recursion now handled via worklist

Algorithm Difference

Before (immediate):

eg.merge(a, b)  # Congruence restored recursively before return
assert eg.find(ax) == eg.find(bx)  # Already equal

After (deferred):

eg.merge(a, b)  # Adds to worklist, returns immediately
assert eg.find(ax) != eg.find(bx)  # Not yet equal

eg.rebuild()  # Process worklist
assert eg.find(ax) == eg.find(bx)  # Now equal

Enables batching multiple merges before a single rebuild, reducing redundant work. Tests updated to call rebuild() after merge operations.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>使用rebuild优化egg库</issue_title>
<issue_description>使用rebuild优化的egg 算法demo:

# 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()

目前的版本所依据的算法demo :

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

from __future__ import annotations

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

</details>



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

- Fixes USTC-KnowledgeComputingLab/ds#157

<!-- 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).

[Copilot]

Pull request overview

This PR implements egg-style deferred rebuilding for the E-Graph, replacing the immediate upward merging strategy with a worklist-based approach that defers congruence restoration until rebuild() is explicitly called.

Key changes:

• Added worklist tracking to defer congruence restoration work
• Modified merge() to add affected E-classes to the worklist instead of immediately repairing
• Introduced new public rebuild() method to process the worklist and restore congruence
• Refactored repair() (previously _repair()) to remove the internal loop, with recursion now handled via the worklist

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	Core implementation changes: added worklist field, modified merge() to defer rebuilding, added rebuild() method, and refactored repair() to work with the worklist-based approach
egg/tests/test_egraph.py	Updated congruence tests to call rebuild() after merge operations to properly test the deferred rebuilding behavior

---

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

[Copilot]

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().