From 8c9edd98996b391613123f2ac92bc8c324bf0886 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 12:42:21 +0000 Subject: [PATCH 1/4] Initial plan From b6b04cad95d0d54bd43dad958b6cef8d077063c3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 12:48:09 +0000 Subject: [PATCH 2/4] Convert egg package to Cython for performance optimization Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- egg/.gitignore | 2 + egg/apyds_egg/{__init__.py => __init__.pyx} | 187 ++++++++++++++------ egg/pyproject.toml | 4 + egg/setup.py | 22 +++ 4 files changed, 156 insertions(+), 59 deletions(-) create mode 100644 egg/.gitignore rename egg/apyds_egg/{__init__.py => __init__.pyx} (54%) create mode 100644 egg/setup.py diff --git a/egg/.gitignore b/egg/.gitignore new file mode 100644 index 0000000..b5c536d --- /dev/null +++ b/egg/.gitignore @@ -0,0 +1,2 @@ +*.cpp +*.c diff --git a/egg/apyds_egg/__init__.py b/egg/apyds_egg/__init__.pyx similarity index 54% rename from egg/apyds_egg/__init__.py rename to egg/apyds_egg/__init__.pyx index b78df7e..33a2614 100644 --- a/egg/apyds_egg/__init__.py +++ b/egg/apyds_egg/__init__.pyx @@ -1,8 +1,12 @@ +# cython: language_level=3 +# cython: boundscheck=False +# cython: wraparound=False +# cython: cdivision=True + from __future__ import annotations __all__ = ["EClassId", "UnionFind", "ENode", "EGraph"] -from dataclasses import dataclass from typing import NewType, Callable, TypeVar, Generic from collections import defaultdict import apyds @@ -12,13 +16,15 @@ T = TypeVar("T") -class UnionFind(Generic[T]): +cdef class UnionFind: """Union-find data structure for managing disjoint sets.""" - + + cdef dict parent + def __init__(self) -> None: - self.parent: dict[T, T] = {} + self.parent = {} - def find(self, x: T) -> T: + cpdef object find(self, object x): """Find the canonical representative of x's set with path compression. Args: @@ -33,7 +39,7 @@ def find(self, x: T) -> T: self.parent[x] = self.find(self.parent[x]) return self.parent[x] - def union(self, a: T, b: T) -> T: + cpdef object union(self, object a, object b): """Union two sets and return the canonical representative. Args: @@ -43,34 +49,63 @@ def union(self, a: T, b: T) -> T: Returns: The canonical representative of the merged set. """ + cdef object ra, rb ra, rb = self.find(a), self.find(b) if ra != rb: self.parent[rb] = ra return ra -@dataclass(frozen=True) -class ENode: +cdef class ENode: """Node in the E-Graph with an operator and children.""" - - op: str - children: tuple[EClassId, ...] - - def canonicalize(self, find: Callable[[EClassId], EClassId]) -> ENode: + + cdef readonly str op + cdef readonly tuple children + cdef int _hash + cdef bint _hash_computed + + def __init__(self, str op, tuple children): + self.op = op + self.children = children + self._hash_computed = False + + def __hash__(self): + if not self._hash_computed: + self._hash = hash((self.op, self.children)) + self._hash_computed = True + return self._hash + + def __eq__(self, other): + if not isinstance(other, ENode): + return False + return self.op == other.op and self.children == other.children + + def __repr__(self): + return f"ENode({self.op!r}, {self.children!r})" + + def canonicalize(self, object find_func): """Canonicalize children using the find function. Args: - find: Function to find the canonical E-class ID. + find_func: 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)) + cdef tuple canon_children = tuple(find_func(c) for c in self.children) + return ENode(self.op, canon_children) -class EGraph: +cdef class EGraph: """E-Graph for representing equivalence classes of terms.""" - + + cdef int _next_id + cdef dict _hashcons + cdef UnionFind _unionfind + cdef dict _classes + cdef object _parents + cdef set _worklist + def __init__(self) -> None: # 1. 唯一性约束 (hashcons): # ENode -> EClassId 的映射。确保具有相同算子且子项属于相同 E-Class 的节点在内存中是唯一的。 @@ -81,20 +116,45 @@ def __init__(self) -> None: # 4. 逆向传播约束 (parents): # EClassId (代表元) -> Set[(ENode, EClassId)] 的映射。记录哪些父节点依赖于该 E-Class。 # 当两个 E-Class 合并时，必须通过此字段通知并更新所有父节点，以维护全等闭包。 - self.next_id: int = 0 - self.hashcons: dict[ENode, EClassId] = {} - 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: + self._next_id = 0 + self._hashcons = {} + self._unionfind = UnionFind() + self._classes = {} + self._parents = defaultdict(set) + self._worklist = set() + + # Expose internal state for compatibility + @property + def next_id(self): + return self._next_id + + @property + def hashcons(self): + return self._hashcons + + @property + def unionfind(self): + return self._unionfind + + @property + def classes(self): + return self._classes + + @property + def parents(self): + return self._parents + + @property + def worklist(self): + return self._worklist + + cdef object _fresh_id(self): """Generate a fresh E-class ID.""" - eid = EClassId(self.next_id) - self.next_id += 1 - return eid + cdef int eid_val = self._next_id + self._next_id += 1 + return EClassId(eid_val) - def find(self, eclass: EClassId) -> EClassId: + cpdef object find(self, object eclass): """Find the canonical representative of an E-class. Args: @@ -103,9 +163,9 @@ def find(self, eclass: EClassId) -> EClassId: Returns: The canonical E-class ID. """ - return self.unionfind.find(eclass) + return self._unionfind.find(eclass) - def add(self, term: apyds.Term) -> EClassId: + cpdef object add(self, object term): """Add a term to the E-Graph and return its E-class ID. Args: @@ -114,15 +174,18 @@ def add(self, term: apyds.Term) -> EClassId: Returns: The E-class ID for the added term. """ - enode = self._term_to_enode(term) + cdef ENode enode = self._term_to_enode(term) return self._add_enode(enode) - def _term_to_enode(self, term: apyds.Term) -> ENode: + cdef ENode _term_to_enode(self, object term): """Convert an apyds.Term to an ENode.""" - inner = term.term + cdef object inner = term.term + cdef list children + cdef object child_term, child_id + cdef int i if isinstance(inner, apyds.List): - children: list[EClassId] = [] + children = [] for i in range(len(inner)): child_term = inner[i] child_id = self.add(child_term) @@ -131,25 +194,26 @@ def _term_to_enode(self, term: apyds.Term) -> ENode: else: return ENode(str(inner), ()) - def _add_enode(self, enode: ENode) -> EClassId: + cdef object _add_enode(self, ENode enode): """Add an ENode to the E-Graph.""" + cdef object eid, c enode = enode.canonicalize(self.find) - if enode in self.hashcons: - return self.find(self.hashcons[enode]) + if enode in self._hashcons: + return self.find(self._hashcons[enode]) eid = self._fresh_id() - self.hashcons[enode] = eid - self.unionfind.parent[eid] = eid - self.classes[eid] = {enode} + self._hashcons[enode] = eid + self._unionfind.parent[eid] = eid + self._classes[eid] = {enode} for c in enode.children: - self.parents[c].add((enode, eid)) + self._parents[c].add((enode, eid)) return eid - def merge(self, a: EClassId, b: EClassId) -> EClassId: + cpdef object merge(self, object a, object b): """Merge two E-classes and defer congruence restoration. Args: @@ -159,23 +223,24 @@ def merge(self, a: EClassId, b: EClassId) -> EClassId: Returns: The canonical E-class ID of the merged class. """ + cdef object ra, rb, r ra, rb = self.find(a), self.find(b) if ra == rb: return ra - r = self.unionfind.union(ra, rb) + r = self._unionfind.union(ra, rb) - self.classes[r] |= self.classes[rb] - del self.classes[rb] + self._classes[r] |= self._classes[rb] + del self._classes[rb] - self.parents[r] |= self.parents[rb] - del self.parents[rb] + self._parents[r] |= self._parents[rb] + del self._parents[rb] - self.worklist.add(r) + self._worklist.add(r) return r - def rebuild(self) -> None: + cpdef void rebuild(self): """Restore congruence by processing the worklist. This method implements the egg-style deferred rebuilding: @@ -183,14 +248,17 @@ def rebuild(self) -> None: - 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() + cdef set todo + cdef object eclass + + while self._worklist: + todo = {self.find(e) for e in self._worklist} + self._worklist.clear() for eclass in todo: self._repair(eclass) - def _repair(self, eclass: EClassId) -> None: + cdef void _repair(self, object eclass): """Restore congruence for a single E-class. This method implements the egg-style repair algorithm: @@ -198,10 +266,11 @@ def _repair(self, eclass: EClassId) -> None: - Merge congruent parents (which may add more work to worklist) - Update hashcons and parent tracking """ - new_parents: dict[ENode, EClassId] = {} + cdef dict new_parents = {} + cdef object pnode, peclass, canon - 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) @@ -210,6 +279,6 @@ def _repair(self, eclass: EClassId) -> None: self.merge(peclass, new_parents[canon]) else: new_parents[canon] = peclass - self.hashcons[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()} diff --git a/egg/pyproject.toml b/egg/pyproject.toml index a62b8ee..35a9bb9 100644 --- a/egg/pyproject.toml +++ b/egg/pyproject.toml @@ -2,6 +2,7 @@ requires = [ "setuptools~=80.9.0", "setuptools-scm~=9.2.2", + "cython~=3.0.11", ] build-backend = "setuptools.build_meta" @@ -28,6 +29,9 @@ root = ".." [tool.setuptools] packages = ["apyds_egg"] +[tool.setuptools.package-data] +apyds_egg = ["*.pyx", "*.pxd", "py.typed"] + [project.optional-dependencies] dev = [ "ruff~=0.14.10", diff --git a/egg/setup.py b/egg/setup.py new file mode 100644 index 0000000..12450b8 --- /dev/null +++ b/egg/setup.py @@ -0,0 +1,22 @@ +from setuptools import setup, Extension +from Cython.Build import cythonize + +extensions = [ + Extension( + "apyds_egg.__init__", + ["apyds_egg/__init__.pyx"], + language="c++", + ) +] + +setup( + ext_modules=cythonize( + extensions, + compiler_directives={ + 'language_level': "3", + 'boundscheck': False, + 'wraparound': False, + 'cdivision': True, + } + ) +) From d480287a47919d1f2833b4a7d3911990f024d364 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 12:51:00 +0000 Subject: [PATCH 3/4] Translate Chinese comments to English for better accessibility Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- egg/apyds_egg/__init__.pyx | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/egg/apyds_egg/__init__.pyx b/egg/apyds_egg/__init__.pyx index 33a2614..518a82f 100644 --- a/egg/apyds_egg/__init__.pyx +++ b/egg/apyds_egg/__init__.pyx @@ -107,15 +107,19 @@ cdef class EGraph: cdef set _worklist def __init__(self) -> None: - # 1. 唯一性约束 (hashcons): - # ENode -> EClassId 的映射。确保具有相同算子且子项属于相同 E-Class 的节点在内存中是唯一的。 - # 2. 等价性维护 (unionfind): - # 管理 EClassId 之间的并查集关系。通过 find 操作将逻辑上的多个 E-Class 映射到唯一的代表元。 - # 3. 集合成员约束 (classes): - # EClassId (代表元) -> Set[ENode] 的映射。存储当前等价类中包含的所有等价项。 - # 4. 逆向传播约束 (parents): - # EClassId (代表元) -> Set[(ENode, EClassId)] 的映射。记录哪些父节点依赖于该 E-Class。 - # 当两个 E-Class 合并时，必须通过此字段通知并更新所有父节点，以维护全等闭包。 + # 1. Uniqueness constraint (hashcons): + # ENode -> EClassId mapping. Ensures nodes with the same operator and children + # belonging to the same E-Class are unique in memory. + # 2. Equivalence maintenance (unionfind): + # Manages union-find relationships between EClassIds. Maps multiple logical + # E-Classes to a unique representative through the find operation. + # 3. Set membership constraint (classes): + # EClassId (representative) -> Set[ENode] mapping. Stores all equivalent terms + # in the current equivalence class. + # 4. Reverse propagation constraint (parents): + # EClassId (representative) -> Set[(ENode, EClassId)] mapping. Records which + # parent nodes depend on this E-Class. When two E-Classes merge, must notify + # and update all parent nodes through this field to maintain congruence closure. self._next_id = 0 self._hashcons = {} self._unionfind = UnionFind() From b411aa77b7c5f8c263a4e9a97ed49c63dd51a771 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 22 Dec 2025 12:52:50 +0000 Subject: [PATCH 4/4] Update README to document Cython optimization Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com> --- egg/README.md | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/egg/README.md b/egg/README.md index 5bf1f1d..18cb9de 100644 --- a/egg/README.md +++ b/egg/README.md @@ -12,6 +12,7 @@ This package implements the egg-style E-Graph data structure with deferred congr - **Deferred Rebuilding**: egg-style deferred rebuilding for performance - **Python Integration**: Seamless integration with apyds terms - **Type-Safe**: Full type hints for Python 3.11+ +- **Cython-Optimized**: Core implementation compiled with Cython for improved performance ## Installation @@ -106,6 +107,7 @@ assert eg.find(fa) == eg.find(fb) - Python 3.11-3.14 - apyds package +- Cython (for building from source) ### Python Package @@ -115,7 +117,7 @@ cd egg # Install dependencies uv sync --extra dev -# Build package +# Build package (compiles Cython extensions) uv build # Run tests @@ -125,6 +127,14 @@ uv run pytest uv run pytest --cov ``` +### Development Notes + +The core implementation is written in Cython (`.pyx` files) for performance optimization. The build process automatically compiles these to C extensions. When developing, you may need to rebuild after changes: + +```bash +pip install -e .[dev] # Reinstall in editable mode after changes +``` + ## License This project is licensed under the GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later).