From 54d4589cb636a28345b94599411d13e5b6dc5cf3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 04:27:46 +0000 Subject: [PATCH 1/3] Initial plan From 22be200e54def87cc77e4f7a47f2d2bacd5409c5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 04:32:53 +0000 Subject: [PATCH 2/3] Fix apyds_egg/__init__.py: make repair private, add are_equal, improve docstrings Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- egg/apyds_egg/__init__.py | 63 ++++++++++++++++++++++++++++++++++----- egg/tests/test_egraph.py | 35 ++++++++++++++++++++++ 2 files changed, 91 insertions(+), 7 deletions(-) diff --git a/egg/apyds_egg/__init__.py b/egg/apyds_egg/__init__.py index 8102323..147d3bc 100644 --- a/egg/apyds_egg/__init__.py +++ b/egg/apyds_egg/__init__.py @@ -17,7 +17,14 @@ def __init__(self) -> None: self.parent: dict[EClassId, EClassId] = {} def find(self, x: EClassId) -> EClassId: - """Find the canonical representative of x's set with path compression.""" + """Find the canonical representative of x's set with path compression. + + Args: + x: The E-class ID to find. + + Returns: + The canonical E-class ID. + """ if x not in self.parent: self.parent[x] = x if self.parent[x] != x: @@ -25,7 +32,15 @@ def find(self, x: EClassId) -> EClassId: return self.parent[x] def union(self, a: EClassId, b: EClassId) -> EClassId: - """Union two sets and return the canonical representative.""" + """Union two sets and return the canonical representative. + + Args: + a: The first E-class ID. + b: The second E-class ID. + + Returns: + The canonical E-class ID of the merged set. + """ ra, rb = self.find(a), self.find(b) if ra != rb: self.parent[rb] = ra @@ -40,7 +55,14 @@ class ENode: children: tuple[EClassId, ...] def canonicalize(self, find: Callable[[EClassId], EClassId]) -> ENode: - """Canonicalize children using the find function.""" + """Canonicalize children using the find function. + + Args: + find: Function to find the canonical E-class ID. + + Returns: + A new ENode with canonicalized children. + """ return ENode(self.op, tuple(find(c) for c in self.children)) @@ -62,7 +84,14 @@ def _fresh_id(self) -> EClassId: return eid def find(self, eclass: EClassId) -> EClassId: - """Find the canonical representative of an E-class.""" + """Find the canonical representative of an E-class. + + Args: + eclass: The E-class ID to find. + + Returns: + The canonical E-class ID. + """ return self.uf.find(eclass) def add(self, term: apyds.Term) -> EClassId: @@ -110,7 +139,15 @@ 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 schedule rebuilding. + + Args: + a: The first E-class ID to merge. + b: The second E-class ID to merge. + + Returns: + The canonical E-class ID of the merged class. + """ ra, rb = self.find(a), self.find(b) if ra == rb: return ra @@ -134,9 +171,21 @@ def rebuild(self) -> None: self.worklist.clear() for eclass in todo: - self.repair(eclass) + self._repair(eclass) + + def are_equal(self, a: EClassId, b: EClassId) -> bool: + """Check if two E-class IDs are equivalent. + + Args: + a: The first E-class ID to compare. + b: The second E-class ID to compare. + + Returns: + True if both E-class IDs belong to the same equivalence class, False otherwise. + """ + return self.find(a) == self.find(b) - def repair(self, eclass: EClassId) -> None: + def _repair(self, eclass: EClassId) -> None: """Repair congruence for an E-class by updating parent nodes.""" new_parents: dict[ENode, EClassId] = {} diff --git a/egg/tests/test_egraph.py b/egg/tests/test_egraph.py index 0f220a0..985d8d1 100644 --- a/egg/tests/test_egraph.py +++ b/egg/tests/test_egraph.py @@ -286,3 +286,38 @@ def test_egraph_hashcons(): assert ab1 == ab2 assert len(eg.classes) == 4 + + +def test_egraph_are_equal(): + eg = EGraph() + + a = eg.add(apyds.Term("a")) + b = eg.add(apyds.Term("b")) + + # Initially they should not be equal + assert not eg.are_equal(a, b) + + # After merging they should be equal + eg.merge(a, b) + assert eg.are_equal(a, b) + + +def test_egraph_are_equal_after_rebuild(): + eg = EGraph() + + x = eg.add(apyds.Term("x")) + 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)")) + + # Initially ax and bx should not be equal + assert not eg.are_equal(ax, bx) + + # Merge a and b + eg.merge(a, b) + eg.rebuild() + + # After rebuild, ax and bx should be equal due to congruence + assert eg.are_equal(ax, bx) From 507b8d7fe709f0bdc401df854d424d4352581956 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 04:39:46 +0000 Subject: [PATCH 3/3] Make UnionFind generic with TypeVar[T] Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- egg/apyds_egg/__init__.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/egg/apyds_egg/__init__.py b/egg/apyds_egg/__init__.py index 147d3bc..c147a06 100644 --- a/egg/apyds_egg/__init__.py +++ b/egg/apyds_egg/__init__.py @@ -3,27 +3,29 @@ __all__ = ["EClassId", "UnionFind", "ENode", "EGraph"] from dataclasses import dataclass -from typing import NewType, Callable +from typing import NewType, Callable, TypeVar, Generic from collections import defaultdict import apyds EClassId = NewType("EClassId", int) +T = TypeVar("T") -class UnionFind: + +class UnionFind(Generic[T]): """Union-find data structure for managing disjoint sets.""" def __init__(self) -> None: - self.parent: dict[EClassId, EClassId] = {} + self.parent: dict[T, T] = {} - def find(self, x: EClassId) -> EClassId: + def find(self, x: T) -> T: """Find the canonical representative of x's set with path compression. Args: - x: The E-class ID to find. + x: The element to find. Returns: - The canonical E-class ID. + The canonical representative of x's set. """ if x not in self.parent: self.parent[x] = x @@ -31,15 +33,15 @@ def find(self, x: EClassId) -> EClassId: self.parent[x] = self.find(self.parent[x]) return self.parent[x] - def union(self, a: EClassId, b: EClassId) -> EClassId: + def union(self, a: T, b: T) -> T: """Union two sets and return the canonical representative. Args: - a: The first E-class ID. - b: The second E-class ID. + a: The first element. + b: The second element. Returns: - The canonical E-class ID of the merged set. + The canonical representative of the merged set. """ ra, rb = self.find(a), self.find(b) if ra != rb: @@ -70,7 +72,7 @@ class EGraph: """E-Graph for representing equivalence classes of terms.""" def __init__(self) -> None: - self.uf = UnionFind() + self.uf = UnionFind[EClassId]() self.next_id = 0 self.classes: dict[EClassId, set[ENode]] = {} self.parents: dict[EClassId, set[tuple[ENode, EClassId]]] = defaultdict(set)