From 0dc2669074a7d2569a856ccec152f1ec790b8f3f Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 06:24:48 +0000
Subject: [PATCH 01/12] Initial plan


From fb7b4722399c0ae8f144e9521b911c41dc4e9dc2 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 06:29:26 +0000
Subject: [PATCH 02/12] Optimize egraph.py execute methods with pre-filtering
 and grouping

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 136 ++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 122 insertions(+), 14 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index c9934b1..fce1ae6 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -55,14 +55,26 @@ def __init__(self) -> None:
         self.terms: set[Term] = set()
         self.facts: set[Term] = set()
         self.pairs: set[Term] = set()
+        # Cache for optimization: maps fact to its equivalent terms
+        self.fact_equiv_cache: dict[Term, set[Term]] = {}
 
     def rebuild(self) -> None:
         self.egraph.rebuild()
+        # Build pairs as before
         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)
+        
+        # Build fact equivalence cache for optimization
+        self.fact_equiv_cache.clear()
+        for fact in self.facts:
+            equiv_terms = set()
+            for term in self.terms:
+                if self.egraph.get_equality(fact, term):
+                    equiv_terms.add(term)
+            self.fact_equiv_cache[fact] = equiv_terms
 
     def add(self, data: Rule) -> None:
         self._add_expr(data)
@@ -84,39 +96,135 @@ def _add_fact(self, data: Rule) -> None:
         self.terms.add(term)
         self.facts.add(term)
 
+    def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) -> list[Term]:
+        """Collect all candidates that can potentially match the pattern."""
+        result = []
+        for candidate in candidates:
+            if pattern @ candidate:
+                result.append(candidate)
+        return result
+
+    def _group_by_equivalence_class(self, terms: list[Term]) -> list[list[Term]]:
+        """Group terms by their equivalence classes in the egraph."""
+        if not terms:
+            return []
+        
+        # Map each term to its representative (canonical form)
+        term_to_repr: dict[Term, EClassId] = {}
+        for term in terms:
+            term_to_repr[term] = self.egraph.core.find(self.egraph._get_or_add(term))
+        
+        # Group terms by representative
+        repr_to_terms: dict[EClassId, list[Term]] = {}
+        for term, repr_id in term_to_repr.items():
+            if repr_id not in repr_to_terms:
+                repr_to_terms[repr_id] = []
+            repr_to_terms[repr_id].append(term)
+        
+        return list(repr_to_terms.values())
+
     def execute(self, data: Rule) -> typing.Iterator[Rule]:
         yield from self._execute_expr(data)
         yield from self._execute_fact(data)
 
     def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
+        """Execute equality query: Q := A == B
+        
+        Optimized approach:
+        1. Collect A_pool = all terms x where A can match x
+        2. Collect B_pool = all terms y where B can match y
+        3. Group both pools by equivalence classes
+        4. For each pair of equivalent groups, check matches
+        """
         lhs_rhs = _extract_lhs_rhs_from_rule(data)
         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)
+            return
+        
+        # Optimization: Collect candidates that can match lhs and rhs
+        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
+        
+        # Group by equivalence classes
+        lhs_groups = self._group_by_equivalence_class(lhs_pool)
+        rhs_groups = self._group_by_equivalence_class(rhs_pool)
+        
+        # Match groups that are equivalent
+        for lhs_group in lhs_groups:
+            for rhs_group in rhs_groups:
+                # Check if these groups are equivalent
+                if lhs_group and rhs_group:
+                    if self.egraph.get_equality(lhs_group[0], rhs_group[0]):
+                        # Try to match within this equivalence class
+                        for x in lhs_group:
+                            for y in rhs_group:
+                                # Build the equality term and check if it matches the query
+                                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]:
+        """Execute fact query: Q := A
+        
+        Optimized approach:
+        1. Collect A_pool = all terms x where A can match x
+        2. For each fact F, use cached equivalent terms (F_pool)
+        3. Group both pools by equivalence classes
+        4. For each pair of equivalent groups, check matches
+        """
         if len(data) != 0:
             return
         idea = data.conclusion
+        
         # 检查是否已经存在严格相等的事实
         for fact in self.facts:
             if self.egraph.get_equality(idea, fact):
                 yield data
-        # 尝试处理含有变量的情况
+                return
+        
+        # Optimization: Collect candidates that can match the idea
+        idea_pool = self._collect_matching_candidates(idea, self.terms)
+        
+        if not idea_pool:
+            return
+        
+        # Group idea_pool by equivalence classes
+        idea_groups = self._group_by_equivalence_class(idea_pool)
+        
+        # For each fact, check if it matches with idea through equivalence
         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])
+            # Get cached equivalent terms for this fact
+            fact_pool = self.fact_equiv_cache.get(fact, set())
+            if not fact_pool:
+                continue
+            
+            # Group fact_pool by equivalence classes
+            fact_groups = self._group_by_equivalence_class(list(fact_pool))
+            
+            # Match groups that are equivalent
+            for idea_group in idea_groups:
+                for fact_group in fact_groups:
+                    # Check if these groups are equivalent
+                    if idea_group and fact_group:
+                        if self.egraph.get_equality(idea_group[0], fact_group[0]):
+                            # Try to match within this equivalence class
+                            for x in idea_group:
+                                for y in fact_group:
+                                    # Build the equality query and check matches
+                                    query = _build_lhs_rhs_to_term(x, y)
+                                    target_equality = _build_lhs_rhs_to_term(idea, fact)
+                                    if unification := query @ target_equality:
+                                        if result := query.ground(unification, scope="1"):
+                                            term = result.term
+                                            if isinstance(term, List):
+                                                yield _build_term_to_rule(term[2])

From fd083e515947d56166acdb9ec36fb2613542aeff Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 06:31:49 +0000
Subject: [PATCH 03/12] Fix early returns and improve documentation in
 optimized execute methods

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index fce1ae6..450f696 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -97,10 +97,21 @@ def _add_fact(self, data: Rule) -> None:
         self.facts.add(term)
 
     def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) -> list[Term]:
-        """Collect all candidates that can potentially match the pattern."""
+        """Collect all candidates that can potentially match the pattern.
+        
+        Uses the @ operator (unification) to check if pattern can match each candidate.
+        The @ operator attempts to unify two terms, returning a substitution if successful.
+        
+        Args:
+            pattern: The pattern term to match against
+            candidates: Set of candidate terms to check
+            
+        Returns:
+            List of candidates that can unify with the pattern
+        """
         result = []
         for candidate in candidates:
-            if pattern @ candidate:
+            if pattern @ candidate:  # Unification check
                 result.append(candidate)
         return result
 
@@ -144,8 +155,8 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
         # 检查是否已经存在严格相等的事实
         if self.egraph.get_equality(lhs, rhs):
             yield data
-            return
         
+        # 尝试处理含有变量的情况
         # Optimization: Collect candidates that can match lhs and rhs
         lhs_pool = self._collect_matching_candidates(lhs, self.terms)
         rhs_pool = self._collect_matching_candidates(rhs, self.terms)
@@ -162,6 +173,7 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
             for rhs_group in rhs_groups:
                 # Check if these groups are equivalent
                 if lhs_group and rhs_group:
+                    # Use the first element as representative of the equivalence class
                     if self.egraph.get_equality(lhs_group[0], rhs_group[0]):
                         # Try to match within this equivalence class
                         for x in lhs_group:
@@ -190,8 +202,9 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
         for fact in self.facts:
             if self.egraph.get_equality(idea, fact):
                 yield data
-                return
+                break
         
+        # 尝试处理含有变量的情况
         # Optimization: Collect candidates that can match the idea
         idea_pool = self._collect_matching_candidates(idea, self.terms)
         
@@ -216,15 +229,17 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
                 for fact_group in fact_groups:
                     # Check if these groups are equivalent
                     if idea_group and fact_group:
+                        # Use the first element as representative of the equivalence class
                         if self.egraph.get_equality(idea_group[0], fact_group[0]):
                             # Try to match within this equivalence class
                             for x in idea_group:
                                 for y in fact_group:
                                     # Build the equality query and check matches
-                                    query = _build_lhs_rhs_to_term(x, y)
-                                    target_equality = _build_lhs_rhs_to_term(idea, fact)
-                                    if unification := query @ target_equality:
-                                        if result := query.ground(unification, scope="1"):
+                                    # query is (x == y), target_pattern is (idea == fact)
+                                    equality_pair = _build_lhs_rhs_to_term(x, y)
+                                    target_pattern = _build_lhs_rhs_to_term(idea, fact)
+                                    if unification := equality_pair @ target_pattern:
+                                        if result := equality_pair.ground(unification, scope="1"):
                                             term = result.term
                                             if isinstance(term, List):
                                                 yield _build_term_to_rule(term[2])

From 921478ac339b16529ce12b85af791656332bac3a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 06:34:08 +0000
Subject: [PATCH 04/12] Remove incorrect break statement in _execute_fact

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index 450f696..ea44534 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -202,7 +202,6 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
         for fact in self.facts:
             if self.egraph.get_equality(idea, fact):
                 yield data
-                break
         
         # 尝试处理含有变量的情况
         # Optimization: Collect candidates that can match the idea

From 54536758a5d85e2d197c7262296565c91801af96 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 06:35:56 +0000
Subject: [PATCH 05/12] Fix variable naming in _execute_fact to match original
 convention

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index ea44534..5dcf724 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -234,11 +234,11 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
                             for x in idea_group:
                                 for y in fact_group:
                                     # Build the equality query and check matches
-                                    # query is (x == y), target_pattern is (idea == fact)
-                                    equality_pair = _build_lhs_rhs_to_term(x, y)
-                                    target_pattern = _build_lhs_rhs_to_term(idea, fact)
-                                    if unification := equality_pair @ target_pattern:
-                                        if result := equality_pair.ground(unification, scope="1"):
+                                    # target is (x == y), query is (idea == fact)
+                                    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])

From 4e14b9f126990de09f3520fe4f695329525fd808 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 07:01:05 +0000
Subject: [PATCH 06/12] Fix formatting to pass pre-commit ruff-format hook

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 42 +++++++++++++++++++++---------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index 5dcf724..df052bc 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -66,7 +66,7 @@ def rebuild(self) -> None:
                 if self.egraph.get_equality(lhs, rhs):
                     equality = _build_lhs_rhs_to_term(lhs, rhs)
                     self.pairs.add(equality)
-        
+
         # Build fact equivalence cache for optimization
         self.fact_equiv_cache.clear()
         for fact in self.facts:
@@ -98,14 +98,14 @@ def _add_fact(self, data: Rule) -> None:
 
     def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) -> list[Term]:
         """Collect all candidates that can potentially match the pattern.
-        
+
         Uses the @ operator (unification) to check if pattern can match each candidate.
         The @ operator attempts to unify two terms, returning a substitution if successful.
-        
+
         Args:
             pattern: The pattern term to match against
             candidates: Set of candidate terms to check
-            
+
         Returns:
             List of candidates that can unify with the pattern
         """
@@ -119,19 +119,19 @@ def _group_by_equivalence_class(self, terms: list[Term]) -> list[list[Term]]:
         """Group terms by their equivalence classes in the egraph."""
         if not terms:
             return []
-        
+
         # Map each term to its representative (canonical form)
         term_to_repr: dict[Term, EClassId] = {}
         for term in terms:
             term_to_repr[term] = self.egraph.core.find(self.egraph._get_or_add(term))
-        
+
         # Group terms by representative
         repr_to_terms: dict[EClassId, list[Term]] = {}
         for term, repr_id in term_to_repr.items():
             if repr_id not in repr_to_terms:
                 repr_to_terms[repr_id] = []
             repr_to_terms[repr_id].append(term)
-        
+
         return list(repr_to_terms.values())
 
     def execute(self, data: Rule) -> typing.Iterator[Rule]:
@@ -140,7 +140,7 @@ def execute(self, data: Rule) -> typing.Iterator[Rule]:
 
     def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
         """Execute equality query: Q := A == B
-        
+
         Optimized approach:
         1. Collect A_pool = all terms x where A can match x
         2. Collect B_pool = all terms y where B can match y
@@ -151,23 +151,23 @@ 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
-        
+
         # 尝试处理含有变量的情况
         # Optimization: Collect candidates that can match lhs and rhs
         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
-        
+
         # Group by equivalence classes
         lhs_groups = self._group_by_equivalence_class(lhs_pool)
         rhs_groups = self._group_by_equivalence_class(rhs_pool)
-        
+
         # Match groups that are equivalent
         for lhs_group in lhs_groups:
             for rhs_group in rhs_groups:
@@ -187,7 +187,7 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
 
     def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
         """Execute fact query: Q := A
-        
+
         Optimized approach:
         1. Collect A_pool = all terms x where A can match x
         2. For each fact F, use cached equivalent terms (F_pool)
@@ -197,32 +197,32 @@ 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
-        
+
         # 尝试处理含有变量的情况
         # Optimization: Collect candidates that can match the idea
         idea_pool = self._collect_matching_candidates(idea, self.terms)
-        
+
         if not idea_pool:
             return
-        
+
         # Group idea_pool by equivalence classes
         idea_groups = self._group_by_equivalence_class(idea_pool)
-        
+
         # For each fact, check if it matches with idea through equivalence
         for fact in self.facts:
             # Get cached equivalent terms for this fact
             fact_pool = self.fact_equiv_cache.get(fact, set())
             if not fact_pool:
                 continue
-            
+
             # Group fact_pool by equivalence classes
             fact_groups = self._group_by_equivalence_class(list(fact_pool))
-            
+
             # Match groups that are equivalent
             for idea_group in idea_groups:
                 for fact_group in fact_groups:

From 56c8ec152d17d1e10cfd85056066113439f79e27 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Tue, 23 Dec 2025 07:26:02 +0000
Subject: [PATCH 07/12] Remove unused self.pairs attribute from Search class

Co-authored-by: hzhangxyz <11623447+hzhangxyz@users.noreply.github.com>
---
 ddss/egraph.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index df052bc..bbcdd55 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -54,18 +54,11 @@ def __init__(self) -> None:
         self.egraph: _EGraph = _EGraph()
         self.terms: set[Term] = set()
         self.facts: set[Term] = set()
-        self.pairs: set[Term] = set()
         # Cache for optimization: maps fact to its equivalent terms
         self.fact_equiv_cache: dict[Term, set[Term]] = {}
 
     def rebuild(self) -> None:
         self.egraph.rebuild()
-        # Build pairs as before
-        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)
 
         # Build fact equivalence cache for optimization
         self.fact_equiv_cache.clear()

From d469b10457d4c0344bec20da09002eb2cb882770 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Tue, 23 Dec 2025 15:34:56 +0800
Subject: [PATCH 08/12] Update egraph.py.

---
 ddss/egraph.py | 110 ++++++++++++++-----------------------------------
 1 file changed, 30 insertions(+), 80 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index bbcdd55..d9dee3c 100644
--- a/ddss/egraph.py
+++ b/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,13 +59,10 @@ def __init__(self) -> None:
         self.egraph: _EGraph = _EGraph()
         self.terms: set[Term] = set()
         self.facts: set[Term] = set()
-        # Cache for optimization: maps fact to its equivalent terms
-        self.fact_equiv_cache: dict[Term, set[Term]] = {}
+        self.fact_equiv_cache: dict[Term, set[Term]] = defaultdict(set)
 
     def rebuild(self) -> None:
         self.egraph.rebuild()
-
-        # Build fact equivalence cache for optimization
         self.fact_equiv_cache.clear()
         for fact in self.facts:
             equiv_terms = set()
@@ -89,57 +91,11 @@ def _add_fact(self, data: Rule) -> None:
         self.terms.add(term)
         self.facts.add(term)
 
-    def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) -> list[Term]:
-        """Collect all candidates that can potentially match the pattern.
-
-        Uses the @ operator (unification) to check if pattern can match each candidate.
-        The @ operator attempts to unify two terms, returning a substitution if successful.
-
-        Args:
-            pattern: The pattern term to match against
-            candidates: Set of candidate terms to check
-
-        Returns:
-            List of candidates that can unify with the pattern
-        """
-        result = []
-        for candidate in candidates:
-            if pattern @ candidate:  # Unification check
-                result.append(candidate)
-        return result
-
-    def _group_by_equivalence_class(self, terms: list[Term]) -> list[list[Term]]:
-        """Group terms by their equivalence classes in the egraph."""
-        if not terms:
-            return []
-
-        # Map each term to its representative (canonical form)
-        term_to_repr: dict[Term, EClassId] = {}
-        for term in terms:
-            term_to_repr[term] = self.egraph.core.find(self.egraph._get_or_add(term))
-
-        # Group terms by representative
-        repr_to_terms: dict[EClassId, list[Term]] = {}
-        for term, repr_id in term_to_repr.items():
-            if repr_id not in repr_to_terms:
-                repr_to_terms[repr_id] = []
-            repr_to_terms[repr_id].append(term)
-
-        return list(repr_to_terms.values())
-
     def execute(self, data: Rule) -> typing.Iterator[Rule]:
         yield from self._execute_expr(data)
         yield from self._execute_fact(data)
 
     def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
-        """Execute equality query: Q := A == B
-
-        Optimized approach:
-        1. Collect A_pool = all terms x where A can match x
-        2. Collect B_pool = all terms y where B can match y
-        3. Group both pools by equivalence classes
-        4. For each pair of equivalent groups, check matches
-        """
         lhs_rhs = _extract_lhs_rhs_from_rule(data)
         if lhs_rhs is None:
             return
@@ -150,28 +106,21 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
             yield data
 
         # 尝试处理含有变量的情况
-        # Optimization: Collect candidates that can match lhs and rhs
         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
 
-        # Group by equivalence classes
         lhs_groups = self._group_by_equivalence_class(lhs_pool)
         rhs_groups = self._group_by_equivalence_class(rhs_pool)
 
-        # Match groups that are equivalent
         for lhs_group in lhs_groups:
             for rhs_group in rhs_groups:
-                # Check if these groups are equivalent
                 if lhs_group and rhs_group:
-                    # Use the first element as representative of the equivalence class
-                    if self.egraph.get_equality(lhs_group[0], rhs_group[0]):
-                        # Try to match within this equivalence class
+                    if self.egraph.get_equality(next(iter(lhs_group)), next(iter(rhs_group))):
                         for x in lhs_group:
                             for y in rhs_group:
-                                # Build the equality term and check if it matches the query
                                 target = _build_lhs_rhs_to_term(x, y)
                                 query = data.conclusion
                                 if unification := target @ query:
@@ -179,14 +128,6 @@ def _execute_expr(self, data: Rule) -> typing.Iterator[Rule]:
                                         yield _build_term_to_rule(result)
 
     def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
-        """Execute fact query: Q := A
-
-        Optimized approach:
-        1. Collect A_pool = all terms x where A can match x
-        2. For each fact F, use cached equivalent terms (F_pool)
-        3. Group both pools by equivalence classes
-        4. For each pair of equivalent groups, check matches
-        """
         if len(data) != 0:
             return
         idea = data.conclusion
@@ -197,37 +138,26 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
                 yield data
 
         # 尝试处理含有变量的情况
-        # Optimization: Collect candidates that can match the idea
         idea_pool = self._collect_matching_candidates(idea, self.terms)
 
         if not idea_pool:
             return
 
-        # Group idea_pool by equivalence classes
         idea_groups = self._group_by_equivalence_class(idea_pool)
 
-        # For each fact, check if it matches with idea through equivalence
         for fact in self.facts:
-            # Get cached equivalent terms for this fact
-            fact_pool = self.fact_equiv_cache.get(fact, set())
+            fact_pool = self.fact_equiv_cache[fact]
             if not fact_pool:
                 continue
 
-            # Group fact_pool by equivalence classes
-            fact_groups = self._group_by_equivalence_class(list(fact_pool))
+            fact_groups = self._group_by_equivalence_class(fact_pool)
 
-            # Match groups that are equivalent
             for idea_group in idea_groups:
                 for fact_group in fact_groups:
-                    # Check if these groups are equivalent
                     if idea_group and fact_group:
-                        # Use the first element as representative of the equivalence class
-                        if self.egraph.get_equality(idea_group[0], fact_group[0]):
-                            # Try to match within this equivalence class
+                        if self.egraph.get_equality(next(iter(idea_group)), next(iter(fact_group))):
                             for x in idea_group:
                                 for y in fact_group:
-                                    # Build the equality query and check matches
-                                    # target is (x == y), query is (idea == fact)
                                     target = _build_lhs_rhs_to_term(x, y)
                                     query = _build_lhs_rhs_to_term(idea, fact)
                                     if unification := target @ query:
@@ -235,3 +165,23 @@ def _execute_fact(self, data: Rule) -> typing.Iterator[Rule]:
                                             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]) -> list[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, list[Term]] = defaultdict(set)
+        for term, eid in term_to_eid.items():
+            eid_to_terms[eid].add(term)
+        return list(eid_to_terms.values())

From 9a26b1fa88c165567ef1172ec3e9548908c5fa50 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hao=20Zhang=28=E5=BC=A0=E6=B5=A9=29?=
 <hzhangxyz@outlook.com>
Date: Tue, 23 Dec 2025 15:53:40 +0800
Subject: [PATCH 09/12] Update ddss/egraph.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 ddss/egraph.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index d9dee3c..429b9a3 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -181,7 +181,7 @@ def _group_by_equivalence_class(self, terms: set[Term]) -> list[set[Term]]:
         for term in terms:
             term_to_eid[term] = self.egraph.find(term)
 
-        eid_to_terms: dict[EClassId, list[Term]] = defaultdict(set)
+        eid_to_terms: dict[EClassId, set[Term]] = defaultdict(set)
         for term, eid in term_to_eid.items():
             eid_to_terms[eid].add(term)
         return list(eid_to_terms.values())

From ab06234a1e2a3e9852994f252b7596236250a76d Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Tue, 23 Dec 2025 15:57:30 +0800
Subject: [PATCH 10/12] Update egraph.py.

---
 ddss/egraph.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index 429b9a3..64d4bd9 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -173,7 +173,7 @@ def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) ->
                 result.add(candidate)
         return result
 
-    def _group_by_equivalence_class(self, terms: set[Term]) -> list[set[Term]]:
+    def _group_by_equivalence_class(self, terms: set[Term]) -> typing.Iterator[set[Term]]:
         if not terms:
             return []
 
@@ -184,4 +184,4 @@ def _group_by_equivalence_class(self, terms: set[Term]) -> list[set[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 list(eid_to_terms.values())
+        return eid_to_terms.values()

From 6ec8d5cb733d77aabec829ca11870f85fe605956 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Tue, 23 Dec 2025 16:05:13 +0800
Subject: [PATCH 11/12] Update egraph.py.

---
 ddss/egraph.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index 64d4bd9..07cd811 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -59,17 +59,26 @@ def __init__(self) -> None:
         self.egraph: _EGraph = _EGraph()
         self.terms: set[Term] = set()
         self.facts: 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()
-        self.fact_equiv_cache.clear()
         for fact in self.facts:
-            equiv_terms = set()
+            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):
-                    equiv_terms.add(term)
-            self.fact_equiv_cache[fact] = equiv_terms
+                    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)
@@ -81,7 +90,9 @@ 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:
@@ -89,7 +100,9 @@ def _add_fact(self, data: Rule) -> None:
             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)

From 3dc85a6a1ed7f8a1ffbc54442f579beaefdf1553 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Tue, 23 Dec 2025 16:06:24 +0800
Subject: [PATCH 12/12] Fix a typing issue.

---
 ddss/egraph.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ddss/egraph.py b/ddss/egraph.py
index 07cd811..50db0dc 100644
--- a/ddss/egraph.py
+++ b/ddss/egraph.py
@@ -186,7 +186,7 @@ def _collect_matching_candidates(self, pattern: Term, candidates: set[Term]) ->
                 result.add(candidate)
         return result
 
-    def _group_by_equivalence_class(self, terms: set[Term]) -> typing.Iterator[set[Term]]:
+    def _group_by_equivalence_class(self, terms: set[Term]) -> typing.Iterable[set[Term]]:
         if not terms:
             return []