Optimize egraph execute methods with pre-filtering and equivalence class grouping

ddss/egraph.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import typing
+from collections import defaultdict
 from apyds import Term, Rule, List
 from apyds_egg import EGraph, EClassId
 
@@ -35,6 +36,10 @@ def _get_or_add(self, data: Term) -> EClassId:
             self.mapping[data] = self.core.add(data)
         return self.mapping[data]
 
+    def find(self, data: Term) -> EClassId:
+        data_id = self._get_or_add(data)
+        return self.core.find(data_id)
+
     def set_equality(self, lhs: Term, rhs: Term) -> None:
         lhs_id = self._get_or_add(lhs)
         rhs_id = self._get_or_add(rhs)
@@ -54,15 +59,26 @@ def __init__(self) -> None:
         self.egraph: _EGraph = _EGraph()
         self.terms: set[Term] = set()
         self.facts: set[Term] = set()
-        self.pairs: set[Term] = set()
+        self.newly_added_terms: set[Term] = set()
+        self.newly_added_facts: set[Term] = set()
+        self.fact_equiv_cache: dict[Term, set[Term]] = defaultdict(set)
 
     def rebuild(self) -> None:
         self.egraph.rebuild()
-        for lhs in self.terms:
-            for rhs in self.terms:
-                if self.egraph.get_equality(lhs, rhs):
-                    equality = _build_lhs_rhs_to_term(lhs, rhs)
-                    self.pairs.add(equality)
+        for fact in self.facts:
+            for term in self.newly_added_terms:
+                if self.egraph.get_equality(fact, term):
+                    self.fact_equiv_cache[fact].add(term)
+        for fact in self.newly_added_facts:
+            for term in self.terms:
+                if self.egraph.get_equality(fact, term):
+                    self.fact_equiv_cache[fact].add(term)
+        for fact in self.newly_added_facts:
+            for term in self.newly_added_terms:
+                if self.egraph.get_equality(fact, term):
+                    self.fact_equiv_cache[fact].add(term)
+        self.newly_added_terms.clear()
+        self.newly_added_facts.clear()
 
     def add(self, data: Rule) -> None:
         self._add_expr(data)
@@ -74,15 +90,19 @@ def _add_expr(self, data: Rule) -> None:
             return
         lhs, rhs = lhs_rhs
         self.terms.add(lhs)
+        self.newly_added_terms.add(lhs)
         self.terms.add(rhs)
+        self.newly_added_terms.add(rhs)
         self.egraph.set_equality(lhs, rhs)
 
     def _add_fact(self, data: Rule) -> None:
         if len(data) != 0:
             return
         term = data.conclusion
         self.terms.add(term)
+        self.newly_added_terms.add(term)
         self.facts.add(term)
+        self.newly_added_facts.add(term)
 
     def execute(self, data: Rule) -> typing.Iterator[Rule]:
         yield from self._execute_expr(data)
@@ -93,30 +113,88 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
         if lhs_rhs is None:
             return
         lhs, rhs = lhs_rhs
+
         # 检查是否已经存在严格相等的事实
         if self.egraph.get_equality(lhs, rhs):
             yield data
+
         # 尝试处理含有变量的情况
-        query = data.conclusion
-        for target in self.pairs:
-            if unification := target @ query:
-                if result := target.ground(unification, scope="1"):
-                    yield _build_term_to_rule(result)
+        lhs_pool = self._collect_matching_candidates(lhs, self.terms)
+        rhs_pool = self._collect_matching_candidates(rhs, self.terms)
+
+        if not lhs_pool or not rhs_pool:
+            return
+
+        lhs_groups = self._group_by_equivalence_class(lhs_pool)
+        rhs_groups = self._group_by_equivalence_class(rhs_pool)
+
+        for lhs_group in lhs_groups:
+            for rhs_group in rhs_groups:
+                if lhs_group and rhs_group:
+                    if self.egraph.get_equality(next(iter(lhs_group)), next(iter(rhs_group))):
+                        for x in lhs_group:
+                            for y in rhs_group:
+                                target = _build_lhs_rhs_to_term(x, y)
+                                query = data.conclusion
+                                if unification := target @ query:
+                                    if result := target.ground(unification, scope="1"):
+                                        yield _build_term_to_rule(result)
 
     def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
         if len(data) != 0:
             return
         idea = data.conclusion
+
         # 检查是否已经存在严格相等的事实
         for fact in self.facts:
             if self.egraph.get_equality(idea, fact):
                 yield data
+
         # 尝试处理含有变量的情况
+        idea_pool = self._collect_matching_candidates(idea, self.terms)
+
+        if not idea_pool:
+            return
+
+        idea_groups = self._group_by_equivalence_class(idea_pool)
+
         for fact in self.facts:
-            query = _build_lhs_rhs_to_term(idea, fact)
-            for target in self.pairs:
-                if unification := target @ query:
-                    if result := target.ground(unification, scope="1"):
-                        term = result.term
-                        if isinstance(term, List):
-                            yield _build_term_to_rule(term[2])
+            fact_pool = self.fact_equiv_cache[fact]
+            if not fact_pool:
+                continue
+
+            fact_groups = self._group_by_equivalence_class(fact_pool)
+
+            for idea_group in idea_groups:
+                for fact_group in fact_groups:
+                    if idea_group and fact_group:
+                        if self.egraph.get_equality(next(iter(idea_group)), next(iter(fact_group))):
+                            for x in idea_group:
+                                for y in fact_group:
+                                    target = _build_lhs_rhs_to_term(x, y)
+                                    query = _build_lhs_rhs_to_term(idea, fact)
+                                    if unification := target @ query:
+                                        if result := target.ground(unification, scope="1"):
+                                            term = result.term
+                                            if isinstance(term, List):
+                                                yield _build_term_to_rule(term[2])
+
+    def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) -> set[Term]:
+        result = set()
+        for candidate in candidates:
+            if pattern @ candidate:
+                result.add(candidate)
+        return result
+
+    def _group_by_equivalence_class(self, terms: set[Term]) -> typing.Iterable[set[Term]]:
+        if not terms:
+            return []
+
+        term_to_eid: dict[Term, EClassId] = {}
+        for term in terms:
+            term_to_eid[term] = self.egraph.find(term)
+
+        eid_to_terms: dict[EClassId, set[Term]] = defaultdict(set)
+        for term, eid in term_to_eid.items():
+            eid_to_terms[eid].add(term)
+        return eid_to_terms.values()