Remove automatic extraction in builtin evaluation

In #265 (released as
9.0.0) the ability to automatically extract a builtin using a global
egraph context was added. This PR removes that feature, requiring all
builtins to be in a normalized form.

I realized that for #241
we want facts to compare structural equality when converting to a boolean,
instead of using the e-graph. Looking at that previous PR, it seems like
a mistake to add this implicit context, making things more confusing
and opaque with minimal UX improvements.

Author: saulshanabrook

Cargo.lock

@@ -4,6 +4,8 @@ _This project uses semantic versioning_
 
 ## UNRELEASED
 
+- Change builtins to not evaluate values in egraph and changes facts to compare structural equality instead of using an egraph when converting to a boolean, removing magic context (`EGraph.current` and `Schedule.current`) that was added in release 9.0.0.
+
 ## 9.0.1 (2025-03-20)
 
 - Add missing i64.log2 method to the bindings

docs/changelog.md

@@ -45,7 +45,7 @@ def _num_rule(a: Num, b: Num, c: Num, i: i64, j: i64):
     yield rewrite(Num(i) * Num(j)).to(Num(i * j))
 
 egraph.saturate()
-egraph.check(eq(expr1).to(expr2))
+egraph.check(expr1 == expr2)
 egraph.extract(expr1)
 ```
 

docs/index.md

@@ -3,6 +3,7 @@
 """
 
 from . import config, ipython_magic  # noqa: F401
+from .bindings import EggSmolError  # noqa: F401
 from .builtins import *  # noqa: UP029
 from .conversion import ConvertError, convert, converter, get_type_args  # noqa: F401
 from .egraph import *

python/egglog/__init__.py

@@ -12,11 +12,9 @@
 
 from typing_extensions import TypeVarTuple, Unpack
 
-from . import bindings
 from .conversion import convert, converter, get_type_args
 from .declarations import *
-from .egraph import BaseExpr, BuiltinExpr, EGraph, expr_fact, function, get_current_ruleset, method
-from .egraph_state import GLOBAL_PY_OBJECT_SORT
+from .egraph import BaseExpr, BuiltinExpr, expr_fact, function, get_current_ruleset, method
 from .functionalize import functionalize
 from .runtime import RuntimeClass, RuntimeExpr, RuntimeFunction
 from .thunk import Thunk
@@ -32,6 +30,7 @@
     "BigRatLike",
     "Bool",
     "BoolLike",
+    "BuiltinEvalError",
     "Map",
     "MapLike",
     "MultiSet",
@@ -56,6 +55,17 @@
 ]
 
 
+class BuiltinEvalError(Exception):
+    """
+    Raised when an builtin cannot be evaluated into a Python primitive because it is complex.
+
+    Try extracting this expression first.
+    """
+
+    def __str__(self) -> str:
+        return f"Cannot evaluate builtin expression into a Python primitive. Try extracting this expression first: {super().__str__()}"
+
+
 class Unit(BuiltinExpr, egg_sort="Unit"):
     """
     The unit type. This is used to reprsent if a value exists in the e-graph or not.
@@ -72,8 +82,8 @@ class String(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> str:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.String)
-        return value.value
+        assert isinstance(value, str)
+        return value
 
     def __init__(self, value: str) -> None: ...
 
@@ -97,8 +107,8 @@ class Bool(BuiltinExpr, egg_sort="bool"):
     @method(preserve=True)
     def eval(self) -> bool:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Bool)
-        return value.value
+        assert isinstance(value, bool)
+        return value
 
     @method(preserve=True)
     def __bool__(self) -> bool:
@@ -132,8 +142,8 @@ class i64(BuiltinExpr):  # noqa: N801
     @method(preserve=True)
     def eval(self) -> int:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Int)
-        return value.value
+        assert isinstance(value, int)
+        return value
 
     @method(preserve=True)
     def __index__(self) -> int:
@@ -251,8 +261,8 @@ class f64(BuiltinExpr):  # noqa: N801
     @method(preserve=True)
     def eval(self) -> float:
         value = _extract_lit(self)
-        assert isinstance(value, bindings.Float)
-        return value.value
+        assert isinstance(value, float)
+        return value
 
     @method(preserve=True)
     def __float__(self) -> float:
@@ -340,9 +350,12 @@ def eval(self) -> dict[T, V]:
         expr = cast("RuntimeExpr", self)
         d = {}
         while call.callable != ClassMethodRef("Map", "empty"):
-            assert call.callable == MethodRef("Map", "insert")
+            msg = "Map can only be evaluated if it is empty or a series of inserts."
+            if call.callable != MethodRef("Map", "insert"):
+                raise BuiltinEvalError(msg)
             call_typed, k_typed, v_typed = call.args
-            assert isinstance(call_typed.expr, CallDecl)
+            if not isinstance(call_typed.expr, CallDecl):
+                raise BuiltinEvalError(msg)
             k = cast("T", expr.__with_expr__(k_typed))
             v = cast("V", expr.__with_expr__(v_typed))
             d[k] = v
@@ -404,7 +417,9 @@ class Set(BuiltinExpr, Generic[T]):
     @method(preserve=True)
     def eval(self) -> set[T]:
         call = _extract_call(self)
-        assert call.callable == InitRef("Set")
+        if call.callable != InitRef("Set"):
+            msg = "Set can only be initialized with the Set constructor."
+            raise BuiltinEvalError(msg)
         return {cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args}
 
     @method(preserve=True)
@@ -466,7 +481,9 @@ class MultiSet(BuiltinExpr, Generic[T]):
     @method(preserve=True)
     def eval(self) -> list[T]:
         call = _extract_call(self)
-        assert call.callable == InitRef("MultiSet")
+        if call.callable != InitRef("MultiSet"):
+            msg = "MultiSet can only be initialized with the MultiSet constructor."
+            raise BuiltinEvalError(msg)
         return [cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args]
 
     @method(preserve=True)
@@ -513,11 +530,15 @@ class Rational(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> Fraction:
         call = _extract_call(self)
-        assert call.callable == InitRef("Rational")
+        if call.callable != InitRef("Rational"):
+            msg = "Rational can only be initialized with the Rational constructor."
+            raise BuiltinEvalError(msg)
 
         def _to_int(e: TypedExprDecl) -> int:
             expr = e.expr
-            assert isinstance(expr, LitDecl)
+            if not isinstance(expr, LitDecl):
+                msg = "Rational can only be initialized with literals"
+                raise BuiltinEvalError(msg)
             assert isinstance(expr.value, int)
             return expr.value
 
@@ -596,9 +617,13 @@ class BigInt(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> int:
         call = _extract_call(self)
-        assert call.callable == ClassMethodRef("BigInt", "from_string")
+        if call.callable != ClassMethodRef("BigInt", "from_string"):
+            msg = "BigInt can only be initialized with the BigInt constructor."
+            raise BuiltinEvalError(msg)
         (s,) = call.args
-        assert isinstance(s.expr, LitDecl)
+        if not isinstance(s.expr, LitDecl):
+            msg = "BigInt can only be initialized with literals"
+            raise BuiltinEvalError(msg)
         assert isinstance(s.expr.value, str)
         return int(s.expr.value)
 
@@ -717,14 +742,19 @@ class BigRat(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> Fraction:
         call = _extract_call(self)
-        assert call.callable == InitRef("BigRat")
+        if call.callable != InitRef("BigRat"):
+            msg = "BigRat can only be initialized with the BigRat constructor."
+            raise BuiltinEvalError(msg)
 
         def _to_fraction(e: TypedExprDecl) -> Fraction:
             expr = e.expr
-            assert isinstance(expr, CallDecl)
-            assert expr.callable == ClassMethodRef("BigInt", "from_string")
+            if not isinstance(expr, CallDecl) or expr.callable != ClassMethodRef("BigInt", "from_string"):
+                msg = "BigRat can only be initialized BigInt strings"
+                raise BuiltinEvalError(msg)
             (s,) = expr.args
-            assert isinstance(s.expr, LitDecl)
+            if not isinstance(s.expr, LitDecl):
+                msg = "BigInt can only be initialized with literals"
+                raise BuiltinEvalError(msg)
             assert isinstance(s.expr.value, str)
             return Fraction(s.expr.value)
 
@@ -821,7 +851,10 @@ def eval(self) -> tuple[T, ...]:
         call = _extract_call(self)
         if call.callable == ClassMethodRef("Vec", "empty"):
             return ()
-        assert call.callable == InitRef("Vec")
+
+        if call.callable != InitRef("Vec"):
+            msg = "Vec can only be initialized with the Vec constructor."
+            raise BuiltinEvalError(msg)
         return tuple(cast("T", cast("RuntimeExpr", self).__with_expr__(x)) for x in call.args)
 
     @method(preserve=True)
@@ -889,10 +922,11 @@ def set(self, index: i64Like, value: T) -> Vec[T]: ...
 class PyObject(BuiltinExpr):
     @method(preserve=True)
     def eval(self) -> object:
-        report = (EGraph.current or EGraph())._run_extract(cast("RuntimeExpr", self), 0)
-        assert isinstance(report, bindings.Best)
-        expr = report.termdag.term_to_expr(report.term, bindings.PanicSpan())
-        return GLOBAL_PY_OBJECT_SORT.load(expr)
+        expr = cast("RuntimeExpr", self).__egg_typed_expr__.expr
+        if not isinstance(expr, PyObjectDecl):
+            msg = "PyObject can only be evaluated if it is a PyObject literal"
+            raise BuiltinEvalError(msg)
+        return expr.value
 
     def __init__(self, value: object) -> None: ...
 
@@ -1027,22 +1061,23 @@ def value_to_annotation(a: object) -> type | None:
 converter(FunctionType, UnstableFn, _convert_function)
 
 
-def _extract_lit(e: BaseExpr) -> bindings._Literal:
+def _extract_lit(e: BaseExpr) -> LitType:
     """
     Special case extracting literals to make this faster by using termdag directly.
     """
-    report = (EGraph.current or EGraph())._run_extract(cast("RuntimeExpr", e), 0)
-    assert isinstance(report, bindings.Best)
-    term = report.term
-    assert isinstance(term, bindings.TermLit)
-    return term.value
+    expr = cast("RuntimeExpr", e).__egg_typed_expr__.expr
+    if not isinstance(expr, LitDecl):
+        msg = "Expected a literal"
+        raise BuiltinEvalError(msg)
+    return expr.value
 
 
 def _extract_call(e: BaseExpr) -> CallDecl:
     """
     Extracts the call form of an expression
     """
-    extracted = cast("RuntimeExpr", (EGraph.current or EGraph()).extract(e))
-    expr = extracted.__egg_typed_expr__.expr
-    assert isinstance(expr, CallDecl)
+    expr = cast("RuntimeExpr", e).__egg_typed_expr__.expr
+    if not isinstance(expr, CallDecl):
+        msg = "Expected a call expression"
+        raise BuiltinEvalError(msg)
     return expr

python/egglog/builtins.py

@@ -4,7 +4,7 @@
 import inspect
 import pathlib
 import tempfile
-from collections.abc import Callable, Generator, Iterable, Iterator
+from collections.abc import Callable, Generator, Iterable
 from contextvars import ContextVar
 from dataclasses import InitVar, dataclass, field
 from functools import partial
@@ -17,7 +17,6 @@
     Generic,
     Literal,
     Never,
-    Protocol,
     TypeAlias,
     TypedDict,
     TypeVar,
@@ -85,7 +84,6 @@
     "set_",
     "simplify",
     "subsume",
-    "try_evaling",
     "union",
     "unstable_combine_rulesets",
     "var",
@@ -847,7 +845,6 @@ class EGraph:
     Can run actions, check facts, run schedules, or extract minimal cost expressions.
     """
 
-    current: ClassVar[EGraph | None] = None
     seminaive: InitVar[bool] = True
     save_egglog_string: InitVar[bool] = False
 
@@ -1200,16 +1197,6 @@ def to_json() -> str:
             if visualize:
                 VisualizerWidget(egraphs=egraphs).display_or_open()
 
-    @contextlib.contextmanager
-    def set_current(self) -> Iterator[None]:
-        """
-        Context manager that will set the current egraph. It will be set back after.
-        """
-        prev_current = EGraph.current
-        EGraph.current = self
-        yield
-        EGraph.current = prev_current
-
     @property
     def _egraph(self) -> bindings.EGraph:
         return self._state.egraph
@@ -1303,8 +1290,6 @@ class Schedule(DelayedDeclerations):
     A composition of some rulesets, either composing them sequentially, running them repeatedly, running them till saturation, or running until some facts are met
     """
 
-    current: ClassVar[Schedule | None] = None
-
     # Defer declerations so that we can have rule generators that used not yet defined yet
     schedule: ScheduleDecl
 
@@ -1332,16 +1317,6 @@ def __add__(self, other: Schedule) -> Schedule:
         """
         return Schedule(Thunk.fn(Declarations.create, self, other), SequenceDecl((self.schedule, other.schedule)))
 
-    @contextlib.contextmanager
-    def set_current(self) -> Iterator[None]:
-        """
-        Context manager that will set the current schedule. It will be set back after
-        """
-        prev_current = Schedule.current
-        Schedule.current = self
-        yield
-        Schedule.current = prev_current
-
 
 @dataclass
 class Ruleset(Schedule):
@@ -1488,9 +1463,12 @@ def __repr__(self) -> str:
 
     def __bool__(self) -> bool:
         """
-        Returns True if the two sides of an equality are equal in the egraph or the expression is in the egraph.
+        Returns True if the two sides of an equality are structurally equal.
         """
-        return (EGraph.current or EGraph()).check_bool(self)
+        if not isinstance(self.fact, EqDecl):
+            msg = "Can only check equality facts"
+            raise TypeError(msg)
+        return self.fact.left == self.fact.right
 
 
 @dataclass
@@ -1876,34 +1854,3 @@ def set_current_ruleset(r: Ruleset | None) -> Generator[None, None, None]:
         yield
     finally:
         _CURRENT_RULESET.reset(token)
-
-
-T_co = TypeVar("T_co", covariant=True)
-
-
-class _EvalsTo(Protocol, Generic[T_co]):
-    def eval(self) -> T_co: ...
-
-
-def try_evaling(schedule: Schedule, expr: Expr, prim_expr: _EvalsTo[T]) -> T:
-    """
-    Try evaling the expression that will result in a primitive expression being fill.
-    if it fails, display the egraph and raise an error.
-    """
-    egraph = EGraph.current or EGraph()
-    with egraph.set_current():
-        try:
-            return prim_expr.eval()
-        except BaseException:  # noqa: S110
-            pass
-    # If this primitive doesn't exist in the egraph, we need to try to create it by
-    # registering the expression and running the schedule
-    egraph.register(expr)
-    egraph.run(Schedule.current or schedule)
-    try:
-        with egraph.set_current():
-            return prim_expr.eval()
-    except BaseException as e:
-        # egraph.display(n_inline_leaves=1, split_primitive_outputs=True)
-        e.add_note(f"Cannot evaluate {egraph.extract(expr)}")
-        raise