Switch to using add_note

Author: saulshanabrook

docs/changelog.md

@@ -11,6 +11,7 @@ _This project uses semantic versioning_
 - Subsumes lambda functions after replacing
 - Add working loopnest test
 - Improve tracebacks on failing conversions.
+- Use `add_note` for exception to add more context, instead of raising a new exception, to make it easier to debug.
 
 ## 8.0.1 (2024-10-24)
 

python/egglog/egraph.py

@@ -42,53 +42,53 @@
 
 
 __all__ = [
+    "Action",
+    "Command",
+    "Command",
     "EGraph",
-    "Module",
-    "function",
-    "ruleset",
-    "method",
-    "relation",
     "Expr",
+    "Fact",
+    "Fact",
+    "GraphvizKwargs",
+    "Module",
+    "RewriteOrRule",
+    "Ruleset",
+    "Schedule",
     "Unit",
-    "rewrite",
+    "_BirewriteBuilder",
+    "_EqBuilder",
+    "_NeBuilder",
+    "_RewriteBuilder",
+    "_SetBuilder",
+    "_UnionBuilder",
+    "action_command",
     "birewrite",
-    "eq",
-    "ne",
-    "panic",
-    "let",
+    "check",
+    "check_eq",
     "constant",
     "delete",
-    "subsume",
-    "union",
-    "set_",
-    "rule",
-    "var",
-    "vars_",
-    "Fact",
-    "expr_parts",
+    "eq",
     "expr_action",
     "expr_fact",
-    "action_command",
-    "Schedule",
+    "expr_parts",
+    "function",
+    "let",
+    "method",
+    "ne",
+    "panic",
+    "relation",
+    "rewrite",
+    "rule",
+    "ruleset",
     "run",
     "seq",
-    "Command",
+    "set_",
     "simplify",
+    "subsume",
+    "union",
     "unstable_combine_rulesets",
-    "check",
-    "GraphvizKwargs",
-    "Ruleset",
-    "_RewriteBuilder",
-    "_BirewriteBuilder",
-    "_EqBuilder",
-    "_NeBuilder",
-    "_SetBuilder",
-    "_UnionBuilder",
-    "RewriteOrRule",
-    "Fact",
-    "Action",
-    "Command",
-    "check_eq",
+    "var",
+    "vars_",
 ]
 
 T = TypeVar("T")
@@ -146,20 +146,29 @@ def simplify(x: EXPR, schedule: Schedule | None = None) -> EXPR:
     return EGraph().extract(x)
 
 
-def check_eq(x: EXPR, y: EXPR, schedule: Schedule | None = None) -> EGraph:
+def check_eq(x: EXPR, y: EXPR, schedule: Schedule | None = None, *, add_second=True, display=False) -> EGraph:
     """
     Verifies that two expressions are equal after running the schedule.
+
+    If add_second is true, then the second expression is added to the egraph before running the schedule.
     """
     egraph = EGraph()
     x_var = egraph.let("__check_eq_x", x)
-    y_var = egraph.let("__check_eq_y", y)
+    y_var: EXPR = egraph.let("__check_eq_y", y) if add_second else y
     if schedule:
-        egraph.run(schedule)
+        try:
+            egraph.run(schedule)
+        finally:
+            if display:
+                egraph.display()
     fact = eq(x_var).to(y_var)
     try:
         egraph.check(fact)
     except bindings.EggSmolError as err:
-        raise AssertionError(f"Failed {eq(x).to(y)}\n -> {ne(egraph.extract(x)).to(egraph.extract(y))})") from err
+        if display:
+            egraph.display()
+        err.add_note(f"Failed:\n{eq(x).to(y)}\n\nExtracted:\n {eq(egraph.extract(x)).to(egraph.extract(y))})")
+        raise
     return egraph
 
 
@@ -598,8 +607,9 @@ def _generate_class_decls(  # noqa: C901,PLR0912
                 unextractable=unextractable,
                 subsume=subsume,
             )
-        except ValueError as e:
-            raise ValueError(f"Error processing {cls_name}.{method_name}: {e}") from e
+        except Exception as e:
+            e.add_note(f"Error processing {cls_name}.{method_name}")
+            raise
 
         if not builtin and not isinstance(ref, InitRef) and not mutates:
             add_default_funcs.append(add_rewrite)
@@ -1389,10 +1399,16 @@ def to_json() -> str:
 
         egraphs = [to_json()]
         i = 0
-        while self.run(schedule or 1).updated and i < max:
-            i += 1
+        # Always visualize, even if we encounter an error
+        try:
+            while (self.run(schedule or 1).updated) and i < max:
+                i += 1
+                egraphs.append(to_json())
+        except:
             egraphs.append(to_json())
-        VisualizerWidget(egraphs=egraphs).display_or_open()
+            raise
+        finally:
+            VisualizerWidget(egraphs=egraphs).display_or_open()
 
     @classmethod
     def current(cls) -> EGraph:

python/egglog/exp/array_api.py

@@ -1671,9 +1671,10 @@ def try_evaling(expr: Expr, prim_expr: i64 | Bool) -> int | bool:
         try:
             expr_extracted = egraph.extract(expr)
         except BaseException as inner_exc:
-            raise ValueError(f"Cannot simplify {expr}") from inner_exc
-        msg = f"Cannot simplify to primitive {expr_extracted}"
-        raise ValueError(msg) from exc
+            inner_exc.add_note(f"Cannot simplify {expr}")
+            raise
+        exc.add_note(f"Cannot simplify to primitive {expr_extracted}")
+        raise
     return egraph.eval(extracted)
 
     # string = (

python/egglog/runtime.py

@@ -29,12 +29,12 @@
 
 __all__ = [
     "LIT_CLASS_NAMES",
-    "resolve_callable",
-    "resolve_type_annotation",
+    "REFLECTED_BINARY_METHODS",
     "RuntimeClass",
     "RuntimeExpr",
     "RuntimeFunction",
-    "REFLECTED_BINARY_METHODS",
+    "resolve_callable",
+    "resolve_type_annotation",
 ]
 
 
@@ -199,7 +199,11 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
         }:
             raise AttributeError
 
-        cls_decl = self.__egg_decls__._classes[self.__egg_tp__.name]
+        try:
+            cls_decl = self.__egg_decls__._classes[self.__egg_tp__.name]
+        except Exception as e:
+            e.add_note(f"Error processing class {self.__egg_tp__.name}")
+            raise
 
         preserved_methods = cls_decl.preserved_methods
         if name in preserved_methods:
@@ -259,7 +263,8 @@ def __call__(self, *args: object, _egg_partial_function: bool = False, **kwargs:
         try:
             signature = self.__egg_decls__.get_callable_decl(self.__egg_ref__).to_function_decl().signature
         except Exception as e:
-            raise TypeError(f"Failed to find callable {self}") from e
+            e.add_note(f"Failed to find callable {self}")
+            raise
         decls = self.__egg_decls__.copy()
         # Special case function application bc we dont support variadic generics yet generally
         if signature == "fn-app":