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/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). diff --git a/egg/apyds_egg/__init__.py b/egg/apyds_egg/__init__.py deleted file mode 100644 index b78df7e..0000000 --- a/egg/apyds_egg/__init__.py +++ /dev/null @@ -1,215 +0,0 @@ -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 - -EClassId = NewType("EClassId", int) - -T = TypeVar("T") - - -class UnionFind(Generic[T]): - """Union-find data structure for managing disjoint sets.""" - - def __init__(self) -> None: - self.parent: dict[T, T] = {} - - def find(self, x: T) -> T: - """Find the canonical representative of x's set with path compression. - - Args: - x: The element to find. - - Returns: - The canonical representative of x's set. - """ - 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: T, b: T) -> T: - """Union two sets and return the canonical representative. - - Args: - a: The first element. - b: The second element. - - Returns: - The canonical representative of the merged set. - """ - ra, rb = self.find(a), self.find(b) - if ra != rb: - self.parent[rb] = ra - return ra - - -@dataclass(frozen=True) -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: - """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)) - - -class EGraph: - """E-Graph for representing equivalence classes of terms.""" - - 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 合并时，必须通过此字段通知并更新所有父节点，以维护全等闭包。 - 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: - """Generate a fresh E-class ID.""" - eid = EClassId(self.next_id) - self.next_id += 1 - return eid - - def find(self, eclass: EClassId) -> EClassId: - """Find the canonical representative of an E-class. - - Args: - eclass: The E-class ID to find. - - Returns: - The canonical E-class ID. - """ - return self.unionfind.find(eclass) - - def add(self, term: apyds.Term) -> EClassId: - """Add a term to the E-Graph and return its E-class ID. - - Args: - term: An apyds.Term to add to the E-Graph. - - Returns: - The E-class ID for the added term. - """ - enode = self._term_to_enode(term) - return self._add_enode(enode) - - def _term_to_enode(self, term: apyds.Term) -> ENode: - """Convert an apyds.Term to an ENode.""" - inner = term.term - - if isinstance(inner, apyds.List): - children: list[EClassId] = [] - for i in range(len(inner)): - child_term = inner[i] - child_id = self.add(child_term) - children.append(child_id) - return ENode("()", tuple(children)) - else: - return ENode(str(inner), ()) - - def _add_enode(self, enode: ENode) -> EClassId: - """Add an ENode to the E-Graph.""" - enode = enode.canonicalize(self.find) - - 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} - - for c in enode.children: - self.parents[c].add((enode, eid)) - - return eid - - def merge(self, a: EClassId, b: EClassId) -> EClassId: - """Merge two E-classes and defer congruence restoration. - - 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 - - r = self.unionfind.union(ra, rb) - - self.classes[r] |= self.classes[rb] - del self.classes[rb] - - self.parents[r] |= self.parents[rb] - del self.parents[rb] - - self.worklist.add(r) - - return r - - def rebuild(self) -> None: - """Restore congruence by processing the worklist. - - 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 (which may add more work to worklist) - - Update hashcons and parent tracking - """ - new_parents: dict[ENode, EClassId] = {} - - for pnode, peclass in list(self.parents[eclass]): - self.hashcons.pop(pnode, None) - - canon = pnode.canonicalize(self.find) - peclass = self.find(peclass) - - if canon in new_parents: - 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()} diff --git a/egg/apyds_egg/__init__.pyx b/egg/apyds_egg/__init__.pyx new file mode 100644 index 0000000..518a82f --- /dev/null +++ b/egg/apyds_egg/__init__.pyx @@ -0,0 +1,288 @@ +# cython: language_level=3 +# cython: boundscheck=False +# cython: wraparound=False +# cython: cdivision=True + +from __future__ import annotations + +__all__ = ["EClassId", "UnionFind", "ENode", "EGraph"] + +from typing import NewType, Callable, TypeVar, Generic +from collections import defaultdict +import apyds + +EClassId = NewType("EClassId", int) + +T = TypeVar("T") + + +cdef class UnionFind: + """Union-find data structure for managing disjoint sets.""" + + cdef dict parent + + def __init__(self) -> None: + self.parent = {} + + cpdef object find(self, object x): + """Find the canonical representative of x's set with path compression. + + Args: + x: The element to find. + + Returns: + The canonical representative of x's set. + """ + 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] + + cpdef object union(self, object a, object b): + """Union two sets and return the canonical representative. + + Args: + a: The first element. + b: The second element. + + 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 + + +cdef class ENode: + """Node in the E-Graph with an operator and children.""" + + 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_func: Function to find the canonical E-class ID. + + Returns: + A new ENode with canonicalized children. + """ + cdef tuple canon_children = tuple(find_func(c) for c in self.children) + return ENode(self.op, canon_children) + + +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. 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() + 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.""" + cdef int eid_val = self._next_id + self._next_id += 1 + return EClassId(eid_val) + + cpdef object find(self, object eclass): + """Find the canonical representative of an E-class. + + Args: + eclass: The E-class ID to find. + + Returns: + The canonical E-class ID. + """ + return self._unionfind.find(eclass) + + cpdef object add(self, object term): + """Add a term to the E-Graph and return its E-class ID. + + Args: + term: An apyds.Term to add to the E-Graph. + + Returns: + The E-class ID for the added term. + """ + cdef ENode enode = self._term_to_enode(term) + return self._add_enode(enode) + + cdef ENode _term_to_enode(self, object term): + """Convert an apyds.Term to an ENode.""" + cdef object inner = term.term + cdef list children + cdef object child_term, child_id + cdef int i + + if isinstance(inner, apyds.List): + children = [] + for i in range(len(inner)): + child_term = inner[i] + child_id = self.add(child_term) + children.append(child_id) + return ENode("()", tuple(children)) + else: + return ENode(str(inner), ()) + + 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]) + + eid = self._fresh_id() + + self._hashcons[enode] = eid + self._unionfind.parent[eid] = eid + self._classes[eid] = {enode} + + for c in enode.children: + self._parents[c].add((enode, eid)) + + return eid + + cpdef object merge(self, object a, object b): + """Merge two E-classes and defer congruence restoration. + + 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. + """ + cdef object ra, rb, r + ra, rb = self.find(a), self.find(b) + if ra == rb: + return ra + + r = self._unionfind.union(ra, rb) + + self._classes[r] |= self._classes[rb] + del self._classes[rb] + + self._parents[r] |= self._parents[rb] + del self._parents[rb] + + self._worklist.add(r) + + return r + + cpdef void rebuild(self): + """Restore congruence by processing the worklist. + + 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 + """ + 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) + + cdef void _repair(self, object eclass): + """Restore congruence for a single E-class. + + This method implements the egg-style repair algorithm: + - Re-canonicalize all parent nodes + - Merge congruent parents (which may add more work to worklist) + - Update hashcons and parent tracking + """ + cdef dict new_parents = {} + cdef object pnode, peclass, canon + + for pnode, peclass in list(self._parents[eclass]): + self._hashcons.pop(pnode, None) + + canon = pnode.canonicalize(self.find) + peclass = self.find(peclass) + + if canon in new_parents: + 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()} 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, + } + ) +)