fix: store CommandDecl

Author: kaeun97

python/egglog/egraph_state.py

@@ -48,6 +48,12 @@ def _normalize_global_let_name(name: str) -> str:
     return name if name.startswith("$") else f"${name}"
 
 
+def _normalize_rule_key(key: str) -> str:
+    """Normalize an egglog rule string for consistent matching."""
+    key = key.replace("'", '"')
+    return re.sub(r"\s+", " ", key).strip()
+
+
 @dataclass
 class EGraphState:
     """
@@ -249,13 +255,12 @@ def _schedule_with_scheduler_to_egg(  # noqa: C901, PLR0912
             case _:
                 assert_never(schedule)
 
-    def translate_rule_key(self, egglog_key: str) -> str:
+    def translate_rule_key(self, egglog_key: str) -> CommandDecl:
         """
-        Translate an egglog rule name to its Python representation.
+        Look up the original Python CommandDecl for an egglog rule key.
         """
-        if egglog_key in self.egg_rule_to_command_decl:
-            return pretty_decl(self.__egg_decls__, self.egg_rule_to_command_decl[egglog_key])
-        return egglog_key
+        normalized = _normalize_rule_key(egglog_key)
+        return self.egg_rule_to_command_decl[normalized]
 
     def ruleset_to_egg(self, ident: Ident) -> None:
         """
@@ -304,7 +309,11 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command
                 else:
                     egg_cmd = bindings.BiRewriteCommand(str(ruleset), rewrite)
 
-                self.egg_rule_to_command_decl[str(egg_cmd)] = cmd
+                normalized = _normalize_rule_key(str(egg_cmd))
+                self.egg_rule_to_command_decl[normalized] = cmd
+                if isinstance(cmd, BiRewriteDecl):
+                    self.egg_rule_to_command_decl[normalized + "=>"] = cmd
+                    self.egg_rule_to_command_decl[normalized + "<="] = cmd
                 return egg_cmd
             case RuleDecl(head, body, name):
                 egg_cmd = bindings.RuleCommand(
@@ -316,7 +325,9 @@ def command_to_egg(self, cmd: CommandDecl, ruleset: Ident) -> bindings._Command
                         str(ruleset),
                     )
                 )
-                self.egg_rule_to_command_decl[str(egg_cmd)] = cmd
+                self.egg_rule_to_command_decl[_normalize_rule_key(str(egg_cmd))] = cmd
+                if name:
+                    self.egg_rule_to_command_decl[name] = cmd
                 return egg_cmd
             # TODO: Replace with just constants value and looking at REF of function
             case DefaultRewriteDecl(ref, expr, subsume):

python/egglog/run_report.py

@@ -1,10 +1,16 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import timedelta
 
 from . import bindings
+from .declarations import CommandDecl, Declarations
 from .egraph_state import EGraphState
+from .pretty import pretty_decl
+
+
+def _format_rule_key(decls: Declarations, key: CommandDecl) -> str:
+    return pretty_decl(decls, key)
 
 
 @dataclass
@@ -25,19 +31,32 @@ def from_bindings(cls, report: bindings.RuleReport) -> RuleReport:
 @dataclass
 class RuleSetReport:
     changed: bool
-    rule_reports: dict[str, list[RuleReport]]
+    rule_reports: dict[CommandDecl, list[RuleReport]]
     search_and_apply_time: timedelta
     merge_time: timedelta
+    _decls: Declarations = field(repr=False, default=None)
 
     @classmethod
-    def from_bindings(cls, report: bindings.RuleSetReport, translate_key: callable) -> RuleSetReport:
+    def from_bindings(
+        cls, report: bindings.RuleSetReport, translate_key: callable, decls: Declarations
+    ) -> RuleSetReport:
         return cls(
             changed=report.changed,
             rule_reports={
                 translate_key(k): [RuleReport.from_bindings(rr) for rr in v] for k, v in report.rule_reports.items()
             },
             search_and_apply_time=report.search_and_apply_time,
             merge_time=report.merge_time,
+            _decls=decls,
+        )
+
+    def __repr__(self) -> str:
+        rule_reports_str = {_format_rule_key(self._decls, k): v for k, v in self.rule_reports.items()}
+        return (
+            f"RuleSetReport(changed={self.changed}, "
+            f"rule_reports={rule_reports_str}, "
+            f"search_and_apply_time={self.search_and_apply_time}, "
+            f"merge_time={self.merge_time})"
         )
 
 
@@ -47,9 +66,11 @@ class IterationReport:
     rebuild_time: timedelta
 
     @classmethod
-    def from_bindings(cls, report: bindings.IterationReport, translate_key: callable) -> IterationReport:
+    def from_bindings(
+        cls, report: bindings.IterationReport, translate_key: callable, decls: Declarations
+    ) -> IterationReport:
         return cls(
-            rule_set_report=RuleSetReport.from_bindings(report.rule_set_report, translate_key),
+            rule_set_report=RuleSetReport.from_bindings(report.rule_set_report, translate_key, decls),
             rebuild_time=report.rebuild_time,
         )
 
@@ -60,26 +81,40 @@ class RunReport:
 
     iterations: list[IterationReport]
     updated: bool
-    search_and_apply_time_per_rule: dict[str, timedelta]
-    num_matches_per_rule: dict[str, int]
+    search_and_apply_time_per_rule: dict[CommandDecl, timedelta]
+    num_matches_per_rule: dict[CommandDecl, int]
     search_and_apply_time_per_ruleset: dict[str, timedelta]
     merge_time_per_ruleset: dict[str, timedelta]
     rebuild_time_per_ruleset: dict[str, timedelta]
+    _decls: Declarations = field(repr=False, default=None)
+
+    def __repr__(self) -> str:
+        time_per_rule = {_format_rule_key(self._decls, k): v for k, v in self.search_and_apply_time_per_rule.items()}
+        matches_per_rule = {_format_rule_key(self._decls, k): v for k, v in self.num_matches_per_rule.items()}
+        return (
+            f"RunReport(iterations={self.iterations}, "
+            f"updated={self.updated}, "
+            f"search_and_apply_time_per_rule={time_per_rule}, "
+            f"num_matches_per_rule={matches_per_rule}, "
+            f"search_and_apply_time_per_ruleset={self.search_and_apply_time_per_ruleset}, "
+            f"merge_time_per_ruleset={self.merge_time_per_ruleset}, "
+            f"rebuild_time_per_ruleset={self.rebuild_time_per_ruleset})"
+        )
 
     @classmethod
     def from_bindings(cls, report: bindings.RunReport, state: EGraphState) -> RunReport:
         return cls(
-            iterations=[IterationReport.from_bindings(it, state.translate_rule_key) for it in report.iterations],
+            iterations=[
+                IterationReport.from_bindings(it, state.translate_rule_key, state.__egg_decls__)
+                for it in report.iterations
+            ],
             updated=report.updated,
             search_and_apply_time_per_rule={
                 state.translate_rule_key(k): v for k, v in report.search_and_apply_time_per_rule.items()
             },
             num_matches_per_rule={state.translate_rule_key(k): v for k, v in report.num_matches_per_rule.items()},
-            search_and_apply_time_per_ruleset={
-                state.translate_rule_key(k): v for k, v in report.search_and_apply_time_per_ruleset.items()
-            },
-            merge_time_per_ruleset={state.translate_rule_key(k): v for k, v in report.merge_time_per_ruleset.items()},
-            rebuild_time_per_ruleset={
-                state.translate_rule_key(k): v for k, v in report.rebuild_time_per_ruleset.items()
-            },
+            search_and_apply_time_per_ruleset=report.search_and_apply_time_per_ruleset,
+            merge_time_per_ruleset=report.merge_time_per_ruleset,
+            rebuild_time_per_ruleset=report.rebuild_time_per_ruleset,
+            _decls=state.__egg_decls__,
         )

python/tests/test_run_report.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 
 from egglog import *
+from egglog.declarations import BiRewriteDecl, RewriteDecl, RuleDecl
 
 
 class TestRunReport:
@@ -35,12 +36,10 @@ def test_rule_names_translated_in_top_level_dicts(self):
         report = egraph.run(10)
 
         for key in report.search_and_apply_time_per_rule:
-            assert "rewrite" in key, f"Expected Python rewrite syntax, got: {key}"
-            assert "__main__" not in key, f"Key should not contain mangled egglog names: {key}"
+            assert isinstance(key, RewriteDecl)
 
         for key in report.num_matches_per_rule:
-            assert "rewrite" in key, f"Expected Python rewrite syntax, got: {key}"
-            assert "__main__" not in key, f"Key should not contain mangled egglog names: {key}"
+            assert isinstance(key, RewriteDecl)
 
     def test_rule_names_translated_in_iterations(self):
         egraph = self._setup_simple_egraph()
@@ -49,8 +48,7 @@ def test_rule_names_translated_in_iterations(self):
         assert len(report.iterations) > 0
         for iteration in report.iterations:
             for key in iteration.rule_set_report.rule_reports:
-                assert "__main__" not in key, f"Iteration rule key not translated: {key}"
-                assert "rewrite" in key, f"Expected Python rewrite syntax, got: {key}"
+                assert isinstance(key, RewriteDecl)
 
     def test_updated_field(self):
         egraph = self._setup_simple_egraph()
@@ -117,7 +115,7 @@ def __mul__(self, other: Math) -> Math: ...
         rule_keys = list(report.search_and_apply_time_per_rule.keys())
         assert len(rule_keys) == 2
         for key in rule_keys:
-            assert "__main__" not in key, f"Key not translated: {key}"
+            assert isinstance(key, RewriteDecl)
 
     def test_empty_run(self):
         egraph = EGraph()
@@ -158,7 +156,7 @@ def __add__(self, other: Num) -> Num: ...
         rule_keys = list(report.search_and_apply_time_per_rule.keys())
         assert len(rule_keys) > 0
         for key in rule_keys:
-            assert "__main__" not in key, f"RuleDecl key not translated: {key}"
+            assert isinstance(key, RuleDecl)
 
     def test_birewrite_decl(self):
         egraph = EGraph()
@@ -178,23 +176,4 @@ def __mul__(self, other: Num) -> Num: ...
         rule_keys = list(report.search_and_apply_time_per_rule.keys())
         assert len(rule_keys) > 0
         for key in rule_keys:
-            assert "__main__" not in key, f"BiRewriteDecl key not translated: {key}"
-            assert "birewrite" in key, f"Expected birewrite() syntax, got: {key}"
-
-    def test_rewrite_decl(self):
-        egraph = EGraph()
-
-        class Num(Expr):
-            def __init__(self, n: i64Like) -> None: ...
-            def __add__(self, other: Num) -> Num: ...
-
-        x, y = vars_("x y", Num)
-        egraph.register(rewrite(x + y).to(y + x))
-        egraph.register(Num(1) + Num(2))
-        report = egraph.run(10)
-
-        rule_keys = list(report.search_and_apply_time_per_rule.keys())
-        assert len(rule_keys) == 1
-        key = rule_keys[0]
-        assert "rewrite" in key, f"Expected rewrite() syntax, got: {key}"
-        assert "__main__" not in key, f"RewriteDecl key not translated: {key}"
+            assert isinstance(key, BiRewriteDecl)