egglog 11.2.0__cp310-cp310-win_amd64.whl → 11.4.0__cp310-cp310-win_amd64.whl

egglog/egraph_state.py

@@ -8,6 +8,7 @@ import re
 from collections import defaultdict
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Literal, overload
+from uuid import UUID
 
 from typing_extensions import assert_never
 
@@ -67,6 +68,7 @@ class EGraphState:
 
     # Bidirectional mapping between egg sort names and python type references.
     type_ref_to_egg_sort: dict[JustTypeRef, str] = field(default_factory=dict)
+    egg_sort_to_type_ref: dict[str, JustTypeRef] = field(default_factory=dict)
 
     # Cache of egg expressions for converting to egg
     expr_to_egg_cache: dict[ExprDecl, bindings._Expr] = field(default_factory=dict)
@@ -85,22 +87,145 @@ class EGraphState:
             egg_fn_to_callable_refs=defaultdict(set, {k: v.copy() for k, v in self.egg_fn_to_callable_refs.items()}),
             callable_ref_to_egg_fn=self.callable_ref_to_egg_fn.copy(),
             type_ref_to_egg_sort=self.type_ref_to_egg_sort.copy(),
+            egg_sort_to_type_ref=self.egg_sort_to_type_ref.copy(),
             expr_to_egg_cache=self.expr_to_egg_cache.copy(),
             cost_callables=self.cost_callables.copy(),
         )
 
-    def schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
+    def run_schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Command:
+        """
+        Turn a run schedule into an egg command.
+
+        If there exists any custom schedulers in the schedule, it will be turned into a custom extract command otherwise
+        will be a normal run command.
+        """
+        processed_schedule = self._process_schedule(schedule)
+        if processed_schedule is None:
+            return bindings.RunSchedule(self._schedule_to_egg(schedule))
+        top_level_schedules = self._schedule_with_scheduler_to_egg(processed_schedule, [])
+        if len(top_level_schedules) == 1:
+            schedule_expr = top_level_schedules[0]
+        else:
+            schedule_expr = bindings.Call(span(), "seq", top_level_schedules)
+        return bindings.UserDefined(span(), "run-schedule", [schedule_expr])
+
+    def _process_schedule(self, schedule: ScheduleDecl) -> ScheduleDecl | None:
+        """
+        Processes a schedule to determine if it contains any custom schedulers.
+
+        If it does, it returns a new schedule with all the required let bindings added to the other scope.
+        If not, returns none.
+
+        Also processes all rulesets in the schedule to make sure they are registered.
+        """
+        bound_schedulers: list[UUID] = []
+        unbound_schedulers: list[BackOffDecl] = []
+
+        def helper(s: ScheduleDecl) -> None:
+            match s:
+                case LetSchedulerDecl(scheduler, inner):
+                    bound_schedulers.append(scheduler.id)
+                    return helper(inner)
+                case RunDecl(ruleset_name, _, scheduler):
+                    self.ruleset_to_egg(ruleset_name)
+                    if scheduler and scheduler.id not in bound_schedulers:
+                        unbound_schedulers.append(scheduler)
+                case SaturateDecl(inner) | RepeatDecl(inner, _):
+                    return helper(inner)
+                case SequenceDecl(schedules):
+                    for sc in schedules:
+                        helper(sc)
+                case _:
+                    assert_never(s)
+            return None
+
+        helper(schedule)
+        if not bound_schedulers and not unbound_schedulers:
+            return None
+        for scheduler in unbound_schedulers:
+            schedule = LetSchedulerDecl(scheduler, schedule)
+        return schedule
+
+    def _schedule_to_egg(self, schedule: ScheduleDecl) -> bindings._Schedule:
+        msg = "Should never reach this, let schedulers should be handled by custom scheduler"
         match schedule:
             case SaturateDecl(schedule):
-                return bindings.Saturate(span(), self.schedule_to_egg(schedule))
+                return bindings.Saturate(span(), self._schedule_to_egg(schedule))
             case RepeatDecl(schedule, times):
-                return bindings.Repeat(span(), times, self.schedule_to_egg(schedule))
+                return bindings.Repeat(span(), times, self._schedule_to_egg(schedule))
             case SequenceDecl(schedules):
-                return bindings.Sequence(span(), [self.schedule_to_egg(s) for s in schedules])
-            case RunDecl(ruleset_name, until):
-                self.ruleset_to_egg(ruleset_name)
+                return bindings.Sequence(span(), [self._schedule_to_egg(s) for s in schedules])
+            case RunDecl(ruleset_name, until, scheduler):
+                if scheduler is not None:
+                    raise ValueError(msg)
                 config = bindings.RunConfig(ruleset_name, None if not until else list(map(self.fact_to_egg, until)))
                 return bindings.Run(span(), config)
+            case LetSchedulerDecl():
+                raise ValueError(msg)
+            case _:
+                assert_never(schedule)
+
+    def _schedule_with_scheduler_to_egg(  # noqa: C901, PLR0912
+        self, schedule: ScheduleDecl, bound_schedulers: list[UUID]
+    ) -> list[bindings._Expr]:
+        """
+        Turns a scheduler into an egg expression, to be used with a custom extract command.
+
+        The bound_schedulers is a list of all the schedulers that have been bound. We can lookup their name as `_scheduler_{index}`.
+        """
+        match schedule:
+            case LetSchedulerDecl(BackOffDecl(id, match_limit, ban_length), inner):
+                name = f"_scheduler_{len(bound_schedulers)}"
+                bound_schedulers.append(id)
+                args: list[bindings._Expr] = []
+                if match_limit is not None:
+                    args.append(bindings.Var(span(), ":match-limit"))
+                    args.append(bindings.Lit(span(), bindings.Int(match_limit)))
+                if ban_length is not None:
+                    args.append(bindings.Var(span(), ":ban-length"))
+                    args.append(bindings.Lit(span(), bindings.Int(ban_length)))
+                back_off_decl = bindings.Call(span(), "back-off", args)
+                let_decl = bindings.Call(span(), "let-scheduler", [bindings.Var(span(), name), back_off_decl])
+                return [let_decl, *self._schedule_with_scheduler_to_egg(inner, bound_schedulers)]
+            case RunDecl(ruleset_name, until, scheduler):
+                args = [bindings.Var(span(), ruleset_name)]
+                if scheduler:
+                    name = "run-with"
+                    scheduler_name = f"_scheduler_{bound_schedulers.index(scheduler.id)}"
+                    args.insert(0, bindings.Var(span(), scheduler_name))
+                else:
+                    name = "run"
+                if until:
+                    if len(until) > 1:
+                        msg = "Can only have one until fact with custom scheduler"
+                        raise ValueError(msg)
+                    args.append(bindings.Var(span(), ":until"))
+                    fact_egg = self.fact_to_egg(until[0])
+                    if isinstance(fact_egg, bindings.Eq):
+                        msg = "Cannot use equality fact with custom scheduler"
+                        raise ValueError(msg)
+                    args.append(fact_egg.expr)
+                return [bindings.Call(span(), name, args)]
+            case SaturateDecl(inner):
+                return [
+                    bindings.Call(span(), "saturate", self._schedule_with_scheduler_to_egg(inner, bound_schedulers))
+                ]
+            case RepeatDecl(inner, times):
+                return [
+                    bindings.Call(
+                        span(),
+                        "repeat",
+                        [
+                            bindings.Lit(span(), bindings.Int(times)),
+                            *self._schedule_with_scheduler_to_egg(inner, bound_schedulers),
+                        ],
+                    )
+                ]
+            case SequenceDecl(schedules):
+                res = []
+                for s in schedules:
+                    res.extend(self._schedule_with_scheduler_to_egg(s, bound_schedulers))
+                return res
             case _:
                 assert_never(schedule)
 
@@ -229,6 +354,7 @@ class EGraphState:
         Creates the egg cost table if needed and gets the name of the table.
         """
         name = self.cost_table_name(ref)
+        print(name, self.cost_callables)
         if ref not in self.cost_callables:
             self.cost_callables.add(ref)
             signature = self.__egg_decls__.get_callable_decl(ref).signature
@@ -332,10 +458,14 @@ class EGraphState:
             pass
         decl = self.__egg_decls__._classes[ref.name]
         self.type_ref_to_egg_sort[ref] = egg_name = decl.egg_name or _generate_type_egg_name(ref)
+        self.egg_sort_to_type_ref[egg_name] = ref
         if not decl.builtin or ref.args:
             if ref.args:
                 if ref.name == "UnstableFn":
                     # UnstableFn is a special case, where the rest of args are collected into a call
+                    if len(ref.args) < 2:
+                        msg = "Zero argument higher order functions not supported"
+                        raise NotImplementedError(msg)
                     type_args: list[bindings._Expr] = [
                         bindings.Call(
                             span(),
@@ -466,11 +596,9 @@ class EGraphState:
                     case _:
                         assert_never(value)
                 res = bindings.Lit(span(), l)
-            case CallDecl(ref, args, _):
-                egg_fn, reverse_args = self.callable_ref_to_egg(ref)
-                egg_args = [self.typed_expr_to_egg(a, False) for a in args]
-                if reverse_args:
-                    egg_args.reverse()
+            case CallDecl() | GetCostDecl():
+                egg_fn, typed_args = self.translate_call(expr_decl)
+                egg_args = [self.typed_expr_to_egg(a, False) for a in typed_args]
                 res = bindings.Call(span(), egg_fn, egg_args)
             case PyObjectDecl(value):
                 res = GLOBAL_PY_OBJECT_SORT.store(value)
@@ -481,11 +609,31 @@ class EGraphState:
                     "unstable-fn",
                     [bindings.Lit(span(), bindings.String(egg_fn_call.name)), *egg_fn_call.args],
                 )
+            case ValueDecl():
+                msg = "Cannot turn a Value into an expression"
+                raise ValueError(msg)
             case _:
                 assert_never(expr_decl.expr)
         self.expr_to_egg_cache[expr_decl] = res
         return res
 
+    def translate_call(self, expr: CallDecl | GetCostDecl) -> tuple[str, list[TypedExprDecl]]:
+        """
+        Handle get cost and call decl, turn into egg table name and typed expr decls.
+        """
+        match expr:
+            case CallDecl(ref, args, _):
+                egg_fn, reverse_args = self.callable_ref_to_egg(ref)
+                args_list = list(args)
+                if reverse_args:
+                    args_list.reverse()
+                return egg_fn, args_list
+            case GetCostDecl(ref, args):
+                cost_table = self.create_cost_table(ref)
+                return cost_table, list(args)
+            case _:
+                assert_never(expr)
+
     def exprs_from_egg(
         self, termdag: bindings.TermDag, terms: list[bindings._Term], tp: JustTypeRef
     ) -> Iterable[TypedExprDecl]:
@@ -529,6 +677,129 @@ class EGraphState:
             case _:
                 assert_never(ref)
 
+    def typed_expr_to_value(self, typed_expr: TypedExprDecl) -> bindings.Value:
+        egg_expr = self.typed_expr_to_egg(typed_expr, False)
+        return self.egraph.eval_expr(egg_expr)[1]
+
+    def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl:  # noqa: C901, PLR0911, PLR0912
+        match tp.name:
+            # Should match list in egraph bindings
+            case "i64":
+                return LitDecl(self.egraph.value_to_i64(value))
+            case "f64":
+                return LitDecl(self.egraph.value_to_f64(value))
+            case "Bool":
+                return LitDecl(self.egraph.value_to_bool(value))
+            case "String":
+                return LitDecl(self.egraph.value_to_string(value))
+            case "Unit":
+                return LitDecl(None)
+            case "PyObject":
+                return PyObjectDecl(self.egraph.value_to_pyobject(GLOBAL_PY_OBJECT_SORT, value))
+            case "Rational":
+                fraction = self.egraph.value_to_rational(value)
+                return CallDecl(
+                    InitRef("Rational"),
+                    (
+                        TypedExprDecl(JustTypeRef("i64"), LitDecl(fraction.numerator)),
+                        TypedExprDecl(JustTypeRef("i64"), LitDecl(fraction.denominator)),
+                    ),
+                )
+            case "BigInt":
+                i = self.egraph.value_to_bigint(value)
+                return CallDecl(
+                    ClassMethodRef("BigInt", "from_string"),
+                    (TypedExprDecl(JustTypeRef("String"), LitDecl(str(i))),),
+                )
+            case "BigRat":
+                fraction = self.egraph.value_to_bigrat(value)
+                return CallDecl(
+                    InitRef("BigRat"),
+                    (
+                        TypedExprDecl(
+                            JustTypeRef("BigInt"),
+                            CallDecl(
+                                ClassMethodRef("BigInt", "from_string"),
+                                (TypedExprDecl(JustTypeRef("String"), LitDecl(str(fraction.numerator))),),
+                            ),
+                        ),
+                        TypedExprDecl(
+                            JustTypeRef("BigInt"),
+                            CallDecl(
+                                ClassMethodRef("BigInt", "from_string"),
+                                (TypedExprDecl(JustTypeRef("String"), LitDecl(str(fraction.denominator))),),
+                            ),
+                        ),
+                    ),
+                )
+            case "Map":
+                k_tp, v_tp = tp.args
+                expr = CallDecl(ClassMethodRef("Map", "empty"), (), (k_tp, v_tp))
+                for k, v in self.egraph.value_to_map(value).items():
+                    expr = CallDecl(
+                        MethodRef("Map", "insert"),
+                        (
+                            TypedExprDecl(tp, expr),
+                            TypedExprDecl(k_tp, self.value_to_expr(k_tp, k)),
+                            TypedExprDecl(v_tp, self.value_to_expr(v_tp, v)),
+                        ),
+                    )
+                return expr
+            case "Set":
+                xs_ = self.egraph.value_to_set(value)
+                (v_tp,) = tp.args
+                return CallDecl(
+                    InitRef("Set"), tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs_), (v_tp,)
+                )
+            case "Vec":
+                xs = self.egraph.value_to_vec(value)
+                (v_tp,) = tp.args
+                return CallDecl(
+                    InitRef("Vec"), tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs), (v_tp,)
+                )
+            case "MultiSet":
+                xs = self.egraph.value_to_multiset(value)
+                (v_tp,) = tp.args
+                return CallDecl(
+                    InitRef("MultiSet"), tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs), (v_tp,)
+                )
+            case "UnstableFn":
+                _names, _args = self.egraph.value_to_function(value)
+                return_tp, *arg_types = tp.args
+                return self._unstable_fn_value_to_expr(_names, _args, return_tp, arg_types)
+        return ValueDecl(value)
+
+    def _unstable_fn_value_to_expr(
+        self, name: str, partial_args: list[bindings.Value], return_tp: JustTypeRef, _arg_types: list[JustTypeRef]
+    ) -> PartialCallDecl:
+        # Similar to FromEggState::from_call but accepts partial list of args and returns in values
+        # Find first callable ref whose return type matches and fill in arg types.
+        for callable_ref in self.egg_fn_to_callable_refs[name]:
+            signature = self.__egg_decls__.get_callable_decl(callable_ref).signature
+            if not isinstance(signature, FunctionSignature):
+                continue
+            if signature.semantic_return_type.name != return_tp.name:
+                continue
+            tcs = TypeConstraintSolver(self.__egg_decls__)
+
+            arg_types, bound_tp_params = tcs.infer_arg_types(
+                signature.arg_types, signature.semantic_return_type, signature.var_arg_type, return_tp, None
+            )
+
+            args = tuple(
+                TypedExprDecl(tp, self.value_to_expr(tp, v)) for tp, v in zip(arg_types, partial_args, strict=False)
+            )
+
+            call_decl = CallDecl(
+                callable_ref,
+                args,
+                # Don't include bound type params if this is just a method, we only needed them for type resolution
+                # but dont need to store them
+                bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else (),
+            )
+            return PartialCallDecl(call_decl)
+        raise ValueError(f"Function '{name}' not found")
+
 
 # https://chatgpt.com/share/9ab899b4-4e17-4426-a3f2-79d67a5ec456
 _EGGLOG_INVALID_IDENT = re.compile(r"[^\w\-+*/?!=<>&|^/%]")
@@ -666,7 +937,7 @@ class FromEggState:
                     args,
                     # Don't include bound type params if this is just a method, we only needed them for type resolution
                     # but dont need to store them
-                    bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else None,
+                    bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else (),
                 )
         raise ValueError(
             f"Could not find callable ref for call {term}. None of these refs matched the types: {self.state.egg_fn_to_callable_refs[term.name]}"

egglog/examples/jointree.py

@@ -62,6 +62,3 @@ egraph.register(query)
 egraph.run(1000)
 print(egraph.extract(query))
 print(egraph.extract(query.size))
-
-
-egraph

egglog/exp/array_api_jit.py

@@ -4,7 +4,7 @@ from typing import TypeVar, cast
 
 import numpy as np
 
-from egglog import EGraph
+from egglog import EGraph, greedy_dag_cost_model
 from egglog.exp.array_api import NDArray, set_array_api_egraph, try_evaling
 from egglog.exp.array_api_numba import array_api_numba_schedule
 from egglog.exp.array_api_program_gen import EvalProgram, array_api_program_gen_schedule, ndarray_function_two_program
@@ -41,7 +41,7 @@ def function_to_program(fn: Callable, save_egglog_string: bool) -> tuple[EGraph,
             res = fn(NDArray.var(arg1), NDArray.var(arg2))
         egraph.register(res)
         egraph.run(array_api_numba_schedule)
-        res_optimized = egraph.extract(res)
+        res_optimized = egraph.extract(res, cost_model=greedy_dag_cost_model())
 
     return (
         egraph,

egglog/pretty.py

@@ -67,7 +67,9 @@ UNARY_METHODS = {
     "__invert__": "~",
 }
 
-AllDecls: TypeAlias = RulesetDecl | CombinedRulesetDecl | CommandDecl | ActionDecl | FactDecl | ExprDecl | ScheduleDecl
+AllDecls: TypeAlias = (
+    RulesetDecl | CombinedRulesetDecl | CommandDecl | ActionDecl | FactDecl | ExprDecl | ScheduleDecl | BackOffDecl
+)
 
 
 def pretty_decl(
@@ -181,17 +183,19 @@ class TraverseContext:
                     if isinstance(de, DefaultRewriteDecl):
                         continue
                     self(de)
-            case CallDecl(ref, exprs, _):
+            case CallDecl(ref, exprs, _) | GetCostDecl(ref, exprs):
                 match ref:
                     case FunctionRef(UnnamedFunctionRef(_, res)):
                         self(res.expr)
                     case _:
                         for e in exprs:
                             self(e.expr)
-            case RunDecl(_, until):
+            case RunDecl(_, until, scheduler):
                 if until:
                     for f in until:
                         self(f)
+                if scheduler:
+                    self(scheduler)
             case PartialCallDecl(c):
                 self(c)
             case CombinedRulesetDecl(_):
@@ -201,6 +205,13 @@ class TraverseContext:
             case SetCostDecl(_, e, c):
                 self(e)
                 self(c)
+            case BackOffDecl() | ValueDecl():
+                pass
+            case LetSchedulerDecl(scheduler, schedule):
+                self(scheduler)
+                self(schedule)
+            case GetCostDecl(ref, args):
+                self(CallDecl(ref, args))
             case _:
                 assert_never(decl)
 
@@ -238,7 +249,11 @@ class PrettyContext:
         # it would take up is > than some constant (~ line length).
         line_diff: int = len(expr) - LINE_DIFFERENCE
         n_parents = self.parents[decl]
-        if n_parents > 1 and n_parents * line_diff > MAX_LINE_LENGTH:
+        if n_parents > 1 and (
+            n_parents * line_diff > MAX_LINE_LENGTH
+            # Schedulers with multiple parents need to be the same object, b/c are created with hidden UUIDs
+            or tp_name == "scheduler"
+        ):
             self.names[decl] = expr_name = self._name_expr(tp_name, expr, copy_identifier=False)
             return expr_name
         return expr
@@ -318,16 +333,31 @@ class PrettyContext:
                     return f"{self(schedules[0], parens=True)} + {self(schedules[1], parens=True)}", "schedule"
                 args = ", ".join(map(self, schedules))
                 return f"seq({args})", "schedule"
-            case RunDecl(ruleset_name, until):
+            case LetSchedulerDecl(scheduler, schedule):
+                return f"{self(scheduler, parens=True)}.scope({self(schedule, parens=True)})", "schedule"
+            case RunDecl(ruleset_name, until, scheduler):
                 ruleset = self.decls._rulesets[ruleset_name]
                 ruleset_str = self(ruleset, ruleset_name=ruleset_name)
-                if not until:
+                if not until and not scheduler:
                     return ruleset_str, "schedule"
-                args = ", ".join(map(self, until))
-                return f"run({ruleset_str}, {args})", "schedule"
+                arg_lst = list(map(self, until or []))
+                if scheduler:
+                    arg_lst.append(f"scheduler={self(scheduler)}")
+                return f"run({ruleset_str}, {', '.join(arg_lst)})", "schedule"
             case DefaultRewriteDecl():
                 msg = "default rewrites should not be pretty printed"
                 raise TypeError(msg)
+            case BackOffDecl(_, match_limit, ban_length):
+                list_args: list[str] = []
+                if match_limit is not None:
+                    list_args.append(f"match_limit={match_limit}")
+                if ban_length is not None:
+                    list_args.append(f"ban_length={ban_length}")
+                return f"back_off({', '.join(list_args)})", "scheduler"
+            case ValueDecl(value):
+                return str(value), "value"
+            case GetCostDecl(ref, args):
+                return f"get_cost({self(CallDecl(ref, args))})", "get_cost"
         assert_never(decl)
 
     def _call(

egglog/runtime.py

@@ -457,7 +457,7 @@ class RuntimeFunction(DelayedDeclerations):
         arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args)
         return_tp = tcs.substitute_typevars(signature.semantic_return_type, cls_name)
         bound_params = (
-            cast("JustTypeRef", bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else None
+            cast("JustTypeRef", bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else ()
         )
         # If we were using unstable-app to call a funciton, add that function back as the first arg.
         if function_value:
@@ -584,11 +584,17 @@ class RuntimeExpr(DelayedDeclerations):
         if (method := _get_expr_method(self, "__eq__")) is not None:
             return method(other)
 
-        # TODO: Check if two objects can be upcasted to be the same. If not, then return NotImplemented so other
-        # expr gets a chance to resolve __eq__ which could be a preserved method.
-        from .egraph import BaseExpr, eq  # noqa: PLC0415
+        if not (isinstance(self, RuntimeExpr) and isinstance(other, RuntimeExpr)):
+            return NotImplemented
+        if self.__egg_typed_expr__.tp != other.__egg_typed_expr__.tp:
+            return NotImplemented
 
-        return eq(cast("BaseExpr", self)).to(cast("BaseExpr", other))
+        from .egraph import Fact  # noqa: PLC0415
+
+        return Fact(
+            Declarations.create(self, other),
+            EqDecl(self.__egg_typed_expr__.tp, self.__egg_typed_expr__.expr, other.__egg_typed_expr__.expr),
+        )
 
     def __ne__(self, other: object) -> object:  # type: ignore[override]
         if (method := _get_expr_method(self, "__ne__")) is not None:
@@ -635,14 +641,22 @@ for name in TYPE_DEFINED_METHODS:
 
 
 for name, r_method in itertools.product(NUMERIC_BINARY_METHODS, (False, True)):
+    method_name = f"__r{name[2:]}" if r_method else name
 
-    def _numeric_binary_method(self: object, other: object, name: str = name, r_method: bool = r_method) -> object:
+    def _numeric_binary_method(
+        self: object, other: object, name: str = name, r_method: bool = r_method, method_name: str = method_name
+    ) -> object:
         """
         Implements numeric binary operations.
 
         Tries to find the minimum cost conversion of either the LHS or the RHS, by finding all methods with either
         the LHS or the RHS as exactly the right type and then upcasting the other to that type.
         """
+        # First check if we have a preserved method for this:
+        if isinstance(self, RuntimeExpr) and (
+            (preserved_method := self.__egg_class_decl__.preserved_methods.get(method_name)) is not None
+        ):
+            return preserved_method.__get__(self)(other)
         # 1. switch if reversed method
         if r_method:
             self, other = other, self
@@ -668,7 +682,6 @@ for name, r_method in itertools.product(NUMERIC_BINARY_METHODS, (False, True)):
         fn = RuntimeFunction(Thunk.value(self.__egg_decls__), Thunk.value(method_ref), self)
         return fn(other)
 
-    method_name = f"__r{name[2:]}" if r_method else name
     setattr(RuntimeExpr, method_name, _numeric_binary_method)
 
 
@@ -688,6 +701,8 @@ def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
             ):
                 raise NotImplementedError(f"Can only turn constants or classvars into callable refs, not {expr}")
             return expr.callable, decl_thunk()
+        case types.MethodWrapperType() if isinstance((slf := callable.__self__), RuntimeClass):
+            return MethodRef(slf.__egg_tp__.name, callable.__name__), slf.__egg_decls__
         case _:
             raise NotImplementedError(f"Cannot turn {callable} of type {type(callable)} into a callable ref")
 

egglog/type_constraint_solver.py

@@ -54,7 +54,7 @@ class TypeConstraintSolver:
         fn_var_args: TypeOrVarRef | None,
         return_: JustTypeRef,
         cls_name: str | None,
-    ) -> tuple[Iterable[JustTypeRef], tuple[JustTypeRef, ...] | None]:
+    ) -> tuple[Iterable[JustTypeRef], tuple[JustTypeRef, ...]]:
         """
         Given a return type, infer the argument types. If there is a variable arg, it returns an infinite iterable.
 
@@ -75,7 +75,7 @@ class TypeConstraintSolver:
                 )
             )
             if cls_name
-            else None
+            else ()
         )
         return arg_types, bound_typevars
 

{egglog-11.2.0.dist-info → egglog-11.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: egglog
-Version: 11.2.0
+Version: 11.4.0
 Classifier: Environment :: MacOS X
 Classifier: Environment :: Win32 (MS Windows)
 Classifier: Intended Audience :: Developers
@@ -51,6 +51,7 @@ Requires-Dist: egglog[array] ; extra == 'docs'
 Requires-Dist: line-profiler ; extra == 'docs'
 Requires-Dist: sphinxcontrib-mermaid ; extra == 'docs'
 Requires-Dist: ablog ; extra == 'docs'
+Requires-Dist: jupytext ; extra == 'docs'
 Provides-Extra: array
 Provides-Extra: dev
 Provides-Extra: test
@@ -71,4 +72,21 @@ allowing you to use e-graphs in Python for optimization, symbolic computation, a
 Please see the [documentation](https://egglog-python.readthedocs.io/) for more information.
 
 Come say hello [on the e-graphs Zulip](https://egraphs.zulipchat.com/#narrow/stream/375765-egglog/) or [open an issue](https://github.com/egraphs-good/egglog-python/issues/new/choose)!
+
+## How to cite
+
+If you use **egglog-python** in academic work, please cite the paper:
+
+```bibtex
+@misc{Shanabrook2023EgglogPython,
+  title         = {Egglog Python: A Pythonic Library for E-graphs},
+  author        = {Saul Shanabrook},
+  year          = {2023},
+  eprint        = {2305.04311},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.PL},
+  doi           = {10.48550/arXiv.2305.04311},
+  url           = {https://arxiv.org/abs/2305.04311},
+  note          = {Presented at EGRAPHS@PLDI 2023}
+}