egglog 6.0.0__cp310-none-win_amd64.whl → 6.1.0__cp310-none-win_amd64.whl

egglog/egraph.py

@@ -170,13 +170,11 @@ class _BaseModule:
 
     @deprecated("Remove this decorator and move the egg_sort to the class statement, i.e. E(Expr, egg_sort='MySort').")
     @overload
-    def class_(self, *, egg_sort: str) -> Callable[[TYPE], TYPE]:
-        ...
+    def class_(self, *, egg_sort: str) -> Callable[[TYPE], TYPE]: ...
 
     @deprecated("Remove this decorator. Simply subclassing Expr is enough now.")
     @overload
-    def class_(self, cls: TYPE, /) -> TYPE:
-        ...
+    def class_(self, cls: TYPE, /) -> TYPE: ...
 
     def class_(self, *args, **kwargs) -> Any:
         """
@@ -201,8 +199,7 @@ class _BaseModule:
         self,
         *,
         preserve: Literal[True],
-    ) -> Callable[[CALLABLE], CALLABLE]:
-        ...
+    ) -> Callable[[CALLABLE], CALLABLE]: ...
 
     @overload
     def method(
@@ -214,8 +211,7 @@ class _BaseModule:
         on_merge: Callable[[Any, Any], Iterable[ActionLike]] | None = None,
         mutates_self: bool = False,
         unextractable: bool = False,
-    ) -> Callable[[CALLABLE], CALLABLE]:
-        ...
+    ) -> Callable[[CALLABLE], CALLABLE]: ...
 
     @overload
     def method(
@@ -228,8 +224,7 @@ class _BaseModule:
         on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
         mutates_self: bool = False,
         unextractable: bool = False,
-    ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]:
-        ...
+    ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
     @deprecated("Use top level method function instead")
     def method(
@@ -249,8 +244,7 @@ class _BaseModule:
         )
 
     @overload
-    def function(self, fn: CALLABLE, /) -> CALLABLE:
-        ...
+    def function(self, fn: CALLABLE, /) -> CALLABLE: ...
 
     @overload
     def function(
@@ -262,8 +256,7 @@ class _BaseModule:
         on_merge: Callable[[Any, Any], Iterable[ActionLike]] | None = None,
         mutates_first_arg: bool = False,
         unextractable: bool = False,
-    ) -> Callable[[CALLABLE], CALLABLE]:
-        ...
+    ) -> Callable[[CALLABLE], CALLABLE]: ...
 
     @overload
     def function(
@@ -276,8 +269,7 @@ class _BaseModule:
         on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
         mutates_first_arg: bool = False,
         unextractable: bool = False,
-    ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]:
-        ...
+    ) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
     @deprecated("Use top level function `function` instead")
     def function(self, *args, **kwargs) -> Any:
@@ -300,24 +292,19 @@ class _BaseModule:
     @overload
     def relation(
         self, name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], tp4: type[E4], /
-    ) -> Callable[[E1, E2, E3, E4], Unit]:
-        ...
+    ) -> Callable[[E1, E2, E3, E4], Unit]: ...
 
     @overload
-    def relation(self, name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], /) -> Callable[[E1, E2, E3], Unit]:
-        ...
+    def relation(self, name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], /) -> Callable[[E1, E2, E3], Unit]: ...
 
     @overload
-    def relation(self, name: str, tp1: type[E1], tp2: type[E2], /) -> Callable[[E1, E2], Unit]:
-        ...
+    def relation(self, name: str, tp1: type[E1], tp2: type[E2], /) -> Callable[[E1, E2], Unit]: ...
 
     @overload
-    def relation(self, name: str, tp1: type[T], /, *, egg_fn: str | None = None) -> Callable[[T], Unit]:
-        ...
+    def relation(self, name: str, tp1: type[T], /, *, egg_fn: str | None = None) -> Callable[[T], Unit]: ...
 
     @overload
-    def relation(self, name: str, /, *, egg_fn: str | None = None) -> Callable[[], Unit]:
-        ...
+    def relation(self, name: str, /, *, egg_fn: str | None = None) -> Callable[[], Unit]: ...
 
     @deprecated("Use top level relation function instead")
     def relation(self, name: str, /, *tps: type, egg_fn: str | None = None) -> Callable[..., Unit]:
@@ -362,8 +349,7 @@ class _BaseModule:
 def method(
     *,
     preserve: Literal[True],
-) -> Callable[[CALLABLE], CALLABLE]:
-    ...
+) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
 # We have to seperate method/function overloads for those that use the T params and those that don't
@@ -380,8 +366,7 @@ def method(
     on_merge: Callable[[Any, Any], Iterable[ActionLike]] | None = None,
     mutates_self: bool = False,
     unextractable: bool = False,
-) -> Callable[[CALLABLE], CALLABLE]:
-    ...
+) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
 @overload
@@ -394,8 +379,7 @@ def method(
     on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
     mutates_self: bool = False,
     unextractable: bool = False,
-) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]:
-    ...
+) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
 
 def method(
@@ -494,11 +478,11 @@ class _ClassDeclerationsConstructor:
         locals = self.frame.f_locals.copy()
         locals[self.cls_name] = self.current_cls
         for k, v in get_type_hints(_Dummytype, globalns=self.frame.f_globals, localns=locals).items():
-            if v.__origin__ == ClassVar:
+            if getattr(v, "__origin__", None) == ClassVar:
                 (inner_tp,) = v.__args__
                 _register_constant(decls, ClassVariableRef(self.cls_name, k), inner_tp, None)
             else:
-                msg = "The only supported annotations on class attributes are class vars"
+                msg = f"On class {self.cls_name}, for attribute '{k}', expected a ClassVar, but got {v}"
                 raise NotImplementedError(msg)
 
         # Then register each of its methods
@@ -565,8 +549,7 @@ class _ClassDeclerationsConstructor:
 
 
 @overload
-def function(fn: CALLABLE, /) -> CALLABLE:
-    ...
+def function(fn: CALLABLE, /) -> CALLABLE: ...
 
 
 @overload
@@ -579,8 +562,7 @@ def function(
     mutates_first_arg: bool = False,
     unextractable: bool = False,
     builtin: bool = False,
-) -> Callable[[CALLABLE], CALLABLE]:
-    ...
+) -> Callable[[CALLABLE], CALLABLE]: ...
 
 
 @overload
@@ -593,8 +575,7 @@ def function(
     on_merge: Callable[[EXPR, EXPR], Iterable[ActionLike]] | None = None,
     mutates_first_arg: bool = False,
     unextractable: bool = False,
-) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]:
-    ...
+) -> Callable[[Callable[P, EXPR]], Callable[P, EXPR]]: ...
 
 
 def function(*args, **kwargs) -> Any:
@@ -760,28 +741,23 @@ def _register_function(
 @overload
 def relation(
     name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], tp4: type[E4], /
-) -> Callable[[E1, E2, E3, E4], Unit]:
-    ...
+) -> Callable[[E1, E2, E3, E4], Unit]: ...
 
 
 @overload
-def relation(name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], /) -> Callable[[E1, E2, E3], Unit]:
-    ...
+def relation(name: str, tp1: type[E1], tp2: type[E2], tp3: type[E3], /) -> Callable[[E1, E2, E3], Unit]: ...
 
 
 @overload
-def relation(name: str, tp1: type[E1], tp2: type[E2], /) -> Callable[[E1, E2], Unit]:
-    ...
+def relation(name: str, tp1: type[E1], tp2: type[E2], /) -> Callable[[E1, E2], Unit]: ...
 
 
 @overload
-def relation(name: str, tp1: type[T], /, *, egg_fn: str | None = None) -> Callable[[T], Unit]:
-    ...
+def relation(name: str, tp1: type[T], /, *, egg_fn: str | None = None) -> Callable[[T], Unit]: ...
 
 
 @overload
-def relation(name: str, /, *, egg_fn: str | None = None) -> Callable[[], Unit]:
-    ...
+def relation(name: str, /, *, egg_fn: str | None = None) -> Callable[[], Unit]: ...
 
 
 def relation(name: str, /, *tps: type, egg_fn: str | None = None) -> Callable[..., Unit]:
@@ -977,7 +953,7 @@ class EGraph(_BaseModule):
         # By default we want to split primitive outputs
         kwargs.setdefault("split_primitive_outputs", True)
         n_inline = kwargs.pop("n_inline_leaves", 0)
-        serialized = self._egraph.serialize(**kwargs)  # type: ignore[misc]
+        serialized = self._egraph.serialize([], **kwargs)  # type: ignore[misc]
         serialized.map_ops(self._state.decls.op_mapping())
         for _ in range(n_inline):
             serialized.inline_leaves()
@@ -1053,12 +1029,10 @@ class EGraph(_BaseModule):
         return cast(EXPR, RuntimeExpr(expr.__egg_decls__, TypedExprDecl(expr.__egg_typed_expr__.tp, VarDecl(name))))
 
     @overload
-    def simplify(self, expr: EXPR, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> EXPR:
-        ...
+    def simplify(self, expr: EXPR, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> EXPR: ...
 
     @overload
-    def simplify(self, expr: EXPR, schedule: Schedule, /) -> EXPR:
-        ...
+    def simplify(self, expr: EXPR, schedule: Schedule, /) -> EXPR: ...
 
     def simplify(
         self, expr: EXPR, limit_or_schedule: int | Schedule, /, *until: Fact, ruleset: Ruleset | None = None
@@ -1102,12 +1076,10 @@ class EGraph(_BaseModule):
         raise NotImplementedError(msg)
 
     @overload
-    def run(self, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> bindings.RunReport:
-        ...
+    def run(self, limit: int, /, *until: Fact, ruleset: Ruleset | None = None) -> bindings.RunReport: ...
 
     @overload
-    def run(self, schedule: Schedule, /) -> bindings.RunReport:
-        ...
+    def run(self, schedule: Schedule, /) -> bindings.RunReport: ...
 
     def run(
         self, limit_or_schedule: int | Schedule, /, *until: Fact, ruleset: Ruleset | None = None
@@ -1147,12 +1119,10 @@ class EGraph(_BaseModule):
         return bindings.Check(egg_facts)
 
     @overload
-    def extract(self, expr: EXPR, /, include_cost: Literal[False] = False) -> EXPR:
-        ...
+    def extract(self, expr: EXPR, /, include_cost: Literal[False] = False) -> EXPR: ...
 
     @overload
-    def extract(self, expr: EXPR, /, include_cost: Literal[True]) -> tuple[EXPR, int]:
-        ...
+    def extract(self, expr: EXPR, /, include_cost: Literal[True]) -> tuple[EXPR, int]: ...
 
     def extract(self, expr: EXPR, include_cost: bool = False) -> EXPR | tuple[EXPR, int]:
         """
@@ -1229,24 +1199,19 @@ class EGraph(_BaseModule):
         self.pop()
 
     @overload
-    def eval(self, expr: i64) -> int:
-        ...
+    def eval(self, expr: i64) -> int: ...
 
     @overload
-    def eval(self, expr: f64) -> float:
-        ...
+    def eval(self, expr: f64) -> float: ...
 
     @overload
-    def eval(self, expr: Bool) -> bool:
-        ...
+    def eval(self, expr: Bool) -> bool: ...
 
     @overload
-    def eval(self, expr: String) -> str:
-        ...
+    def eval(self, expr: String) -> str: ...
 
     @overload
-    def eval(self, expr: PyObject) -> object:
-        ...
+    def eval(self, expr: PyObject) -> object: ...
 
     def eval(self, expr: Expr) -> object:
         """
@@ -1347,8 +1312,7 @@ class Unit(Expr, egg_sort="Unit", builtin=True):
     The unit type. This is also used to reprsent if a value exists, if it is resolved or not.
     """
 
-    def __init__(self) -> None:
-        ...
+    def __init__(self) -> None: ...
 
 
 def ruleset(
@@ -1501,6 +1465,7 @@ class Rewrite(Command):
     _lhs: RuntimeExpr
     _rhs: RuntimeExpr
     _conditions: tuple[Fact, ...]
+    _subsume: bool
     _fn_name: ClassVar[str] = "rewrite"
 
     def __str__(self) -> str:
@@ -1509,7 +1474,7 @@ class Rewrite(Command):
 
     def _to_egg_command(self, default_ruleset_name: str) -> bindings._Command:
         return bindings.RewriteCommand(
-            self.ruleset.egg_name if self.ruleset else default_ruleset_name, self._to_egg_rewrite()
+            self.ruleset.egg_name if self.ruleset else default_ruleset_name, self._to_egg_rewrite(), self._subsume
         )
 
     def _to_egg_rewrite(self) -> bindings.Rewrite:
@@ -1524,7 +1489,7 @@ class Rewrite(Command):
         return Declarations.create(self._lhs, self._rhs, *self._conditions)
 
     def with_ruleset(self, ruleset: Ruleset) -> Rewrite:
-        return Rewrite(ruleset, self._lhs, self._rhs, self._conditions)
+        return Rewrite(ruleset, self._lhs, self._rhs, self._conditions, self._subsume)
 
 
 @dataclass
@@ -1688,21 +1653,23 @@ class ExprAction(Action):
 
 
 @dataclass
-class Delete(Action):
+class Change(Action):
     """
-    Remove a function call from an EGraph.
+    Change a function call in an EGraph.
     """
 
+    change: Literal["delete", "subsume"]
     _call: RuntimeExpr
 
     def __str__(self) -> str:
-        return f"delete({self._call})"
+        return f"{self.change}({self._call})"
 
-    def _to_egg_action(self) -> bindings.Delete:
+    def _to_egg_action(self) -> bindings.Change:
         egg_call = self._call.__egg__
         if not isinstance(egg_call, bindings.Call):
             raise ValueError(f"Can only create a call with a call for the lhs, got {self._call}")  # noqa: TRY004
-        return bindings.Delete(egg_call.name, egg_call.args)
+        change: bindings._Change = bindings.Delete() if self.change == "delete" else bindings.Subsume()
+        return bindings.Change(change, egg_call.name, egg_call.args)
 
     @property
     def __egg_decls__(self) -> Declarations:
@@ -1836,29 +1803,25 @@ class Sequence(Schedule):
 
 @deprecated("Use <ruleset>.register(<rewrite>) instead of passing rulesets as arguments to rewrites.")
 @overload
-def rewrite(lhs: EXPR, ruleset: Ruleset) -> _RewriteBuilder[EXPR]:
-    ...
+def rewrite(lhs: EXPR, ruleset: Ruleset, *, subsume: bool = False) -> _RewriteBuilder[EXPR]: ...
 
 
 @overload
-def rewrite(lhs: EXPR, ruleset: None = None) -> _RewriteBuilder[EXPR]:
-    ...
+def rewrite(lhs: EXPR, ruleset: None = None, *, subsume: bool = False) -> _RewriteBuilder[EXPR]: ...
 
 
-def rewrite(lhs: EXPR, ruleset: Ruleset | None = None) -> _RewriteBuilder[EXPR]:
+def rewrite(lhs: EXPR, ruleset: Ruleset | None = None, *, subsume: bool = False) -> _RewriteBuilder[EXPR]:
     """Rewrite the given expression to a new expression."""
-    return _RewriteBuilder(lhs, ruleset)
+    return _RewriteBuilder(lhs, ruleset, subsume)
 
 
 @deprecated("Use <ruleset>.register(<birewrite>) instead of passing rulesets as arguments to birewrites.")
 @overload
-def birewrite(lhs: EXPR, ruleset: Ruleset) -> _BirewriteBuilder[EXPR]:
-    ...
+def birewrite(lhs: EXPR, ruleset: Ruleset) -> _BirewriteBuilder[EXPR]: ...
 
 
 @overload
-def birewrite(lhs: EXPR, ruleset: None = None) -> _BirewriteBuilder[EXPR]:
-    ...
+def birewrite(lhs: EXPR, ruleset: None = None) -> _BirewriteBuilder[EXPR]: ...
 
 
 def birewrite(lhs: EXPR, ruleset: Ruleset | None = None) -> _BirewriteBuilder[EXPR]:
@@ -1892,7 +1855,12 @@ def expr_action(expr: Expr) -> Action:
 
 def delete(expr: Expr) -> Action:
     """Create a delete expression."""
-    return Delete(to_runtime_expr(expr))
+    return Change("delete", to_runtime_expr(expr))
+
+
+def subsume(expr: Expr) -> Action:
+    """Subsume an expression."""
+    return Change("subsume", to_runtime_expr(expr))
 
 
 def expr_fact(expr: Expr) -> Fact:
@@ -1911,13 +1879,11 @@ def set_(lhs: EXPR) -> _SetBuilder[EXPR]:
 
 @deprecated("Use <ruleset>.register(<rule>) instead of passing rulesets as arguments to rules.")
 @overload
-def rule(*facts: FactLike, ruleset: Ruleset, name: str | None = None) -> _RuleBuilder:
-    ...
+def rule(*facts: FactLike, ruleset: Ruleset, name: str | None = None) -> _RuleBuilder: ...
 
 
 @overload
-def rule(*facts: FactLike, ruleset: None = None, name: str | None = None) -> _RuleBuilder:
-    ...
+def rule(*facts: FactLike, ruleset: None = None, name: str | None = None) -> _RuleBuilder: ...
 
 
 def rule(*facts: FactLike, ruleset: Ruleset | None = None, name: str | None = None) -> _RuleBuilder:
@@ -1947,10 +1913,11 @@ def vars_(names: str, bound: type[EXPR]) -> Iterable[EXPR]:
 class _RewriteBuilder(Generic[EXPR]):
     lhs: EXPR
     ruleset: Ruleset | None
+    subsume: bool
 
     def to(self, rhs: EXPR, *conditions: FactLike) -> Rewrite:
         lhs = to_runtime_expr(self.lhs)
-        rule = Rewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions))
+        rule = Rewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions), self.subsume)
         if self.ruleset:
             self.ruleset.append(rule)
         return rule
@@ -1966,7 +1933,7 @@ class _BirewriteBuilder(Generic[EXPR]):
 
     def to(self, rhs: EXPR, *conditions: FactLike) -> Command:
         lhs = to_runtime_expr(self.lhs)
-        rule = BiRewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions))
+        rule = BiRewrite(self.ruleset, lhs, convert_to_same_type(rhs, lhs), _fact_likes(conditions), False)
         if self.ruleset:
             self.ruleset.append(rule)
         return rule

egglog/examples/bool.py

@@ -27,8 +27,7 @@ R = relation("R", i64)
 
 
 @function
-def f(i: i64Like) -> Bool:
-    ...
+def f(i: i64Like) -> Bool: ...
 
 
 i = var("i", i64)

egglog/examples/eqsat_basic.py

@@ -3,6 +3,7 @@
 Basic equality saturation example.
 ==================================
 """
+
 from __future__ import annotations
 
 from egglog import *
@@ -11,18 +12,14 @@ egraph = EGraph()
 
 
 class Num(Expr):
-    def __init__(self, value: i64Like) -> None:
-        ...
+    def __init__(self, value: i64Like) -> None: ...
 
     @classmethod
-    def var(cls, name: StringLike) -> Num:
-        ...
+    def var(cls, name: StringLike) -> Num: ...
 
-    def __add__(self, other: Num) -> Num:
-        ...
+    def __add__(self, other: Num) -> Num: ...
 
-    def __mul__(self, other: Num) -> Num:
-        ...
+    def __mul__(self, other: Num) -> Num: ...
 
 
 expr1 = Num(2) * (Num.var("x") + Num(3))

egglog/examples/fib.py

@@ -3,14 +3,14 @@
 Fibonacci numbers example
 =========================
 """
+
 from __future__ import annotations
 
 from egglog import *
 
 
 @function
-def fib(x: i64Like) -> i64:
-    ...
+def fib(x: i64Like) -> i64: ...
 
 
 f0, f1, x = vars_("f0 f1 x", i64)

egglog/examples/lambda_.py

@@ -4,6 +4,7 @@
 Lambda Calculus
 ===============
 """
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
@@ -22,66 +23,53 @@ class Val(Expr):
     TRUE: ClassVar[Val]
     FALSE: ClassVar[Val]
 
-    def __init__(self, v: i64Like) -> None:
-        ...
+    def __init__(self, v: i64Like) -> None: ...
 
 
 class Var(Expr):
-    def __init__(self, v: StringLike) -> None:
-        ...
+    def __init__(self, v: StringLike) -> None: ...
 
 
 class Term(Expr):
     @classmethod
-    def val(cls, v: Val) -> Term:
-        ...
+    def val(cls, v: Val) -> Term: ...
 
     @classmethod
-    def var(cls, v: Var) -> Term:
-        ...
+    def var(cls, v: Var) -> Term: ...
 
-    def __add__(self, other: Term) -> Term:
-        ...
+    def __add__(self, other: Term) -> Term: ...
 
     def __eq__(self, other: Term) -> Term:  # type: ignore[override]
         ...
 
-    def __call__(self, other: Term) -> Term:
-        ...
+    def __call__(self, other: Term) -> Term: ...
 
-    def eval(self) -> Val:
-        ...
+    def eval(self) -> Val: ...
 
-    def v(self) -> Var:
-        ...
+    def v(self) -> Var: ...
 
 
 @function
-def lam(x: Var, t: Term) -> Term:
-    ...
+def lam(x: Var, t: Term) -> Term: ...
 
 
 @function
-def let_(x: Var, t: Term, b: Term) -> Term:
-    ...
+def let_(x: Var, t: Term, b: Term) -> Term: ...
 
 
 @function
-def fix(x: Var, t: Term) -> Term:
-    ...
+def fix(x: Var, t: Term) -> Term: ...
 
 
 @function
-def if_(c: Term, t: Term, f: Term) -> Term:
-    ...
+def if_(c: Term, t: Term, f: Term) -> Term: ...
 
 
 StringSet = Set[Var]
 
 
 @function(merge=lambda old, new: old & new)
-def freer(t: Term) -> StringSet:
-    ...
+def freer(t: Term) -> StringSet: ...
 
 
 (v, v1, v2) = vars_("v v1 v2", Val)

egglog/examples/matrix.py

@@ -2,6 +2,7 @@
 Matrix multiplication and Kronecker product.
 ============================================
 """
+
 from __future__ import annotations
 
 from egglog import *
@@ -18,8 +19,7 @@ class Dim(Expr):
     """
 
     @method(egg_fn="Lit")
-    def __init__(self, value: i64Like) -> None:
-        ...
+    def __init__(self, value: i64Like) -> None: ...
 
     @method(egg_fn="NamedDim")
     @classmethod