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

egglog/__init__.py

@@ -4,7 +4,7 @@ Package for creating e-graphs in Python.
 
 from . import config, ipython_magic  # noqa: F401
 from .builtins import *  # noqa: UP029
+from .conversion import convert, converter  # noqa: F401
 from .egraph import *
-from .runtime import convert, converter  # noqa: F401
 
 del ipython_magic

egglog/bindings.pyi

@@ -29,7 +29,9 @@ class EGraph:
         fact_directory: str | Path | None = None,
         seminaive: bool = True,
         terms_encoding: bool = False,
+        record: bool = False,
     ) -> None: ...
+    def commands(self) -> str | None: ...
     def parse_program(self, __input: str, /) -> list[_Command]: ...
     def run_program(self, *commands: _Command) -> list[str]: ...
     def extract_report(self) -> _ExtractReport | None: ...
@@ -492,6 +494,12 @@ class Relation:
 class PrintOverallStatistics:
     def __init__(self) -> None: ...
 
+@final
+class UnstableCombinedRuleset:
+    name: str
+    rulesets: list[str]
+    def __init__(self, name: str, rulesets: list[str]) -> None: ...
+
 _Command: TypeAlias = (
     SetOption
     | Datatype
@@ -519,6 +527,7 @@ _Command: TypeAlias = (
     | CheckProof
     | Relation
     | PrintOverallStatistics
+    | UnstableCombinedRuleset
 )
 
 def termdag_term_to_expr(termdag: TermDag, term: _Term) -> _Expr: ...

egglog/builtins.py

@@ -5,10 +5,14 @@ Builtin sorts and function to egg.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union
+from functools import partial
+from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, overload
 
+from typing_extensions import TypeVarTuple, Unpack
+
+from .conversion import converter
 from .egraph import Expr, Unit, function, method
-from .runtime import converter
+from .runtime import RuntimeFunction
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -31,6 +35,7 @@ __all__ = [
     "py_eval",
     "py_exec",
     "py_eval_fn",
+    "UnstableFn",
 ]
 
 
@@ -461,3 +466,38 @@ def py_exec(code: StringLike, globals: object = PyObject.dict(), locals: object
     """
     Copies the locals, execs the Python code, and returns the locals with any updates.
     """
+
+
+TS = TypeVarTuple("TS")
+
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+
+
+class UnstableFn(Expr, Generic[T, Unpack[TS]], builtin=True):
+    @overload
+    def __init__(self, f: Callable[[Unpack[TS]], T]) -> None: ...
+
+    @overload
+    def __init__(self, f: Callable[[T1, Unpack[TS]], T], _a: T1, /) -> None: ...
+
+    @overload
+    def __init__(self, f: Callable[[T1, T2, Unpack[TS]], T], _a: T1, _b: T2, /) -> None: ...
+
+    # Removing due to bug in MyPy
+    # https://github.com/python/mypy/issues/17212
+    # @overload
+    # def __init__(self, f: Callable[[T1, T2, T3, Unpack[TS]], T], _a: T1, _b: T2, _c: T3, /) -> None: ...
+
+    # etc, for partial application
+
+    @method(egg_fn="unstable-fn")
+    def __init__(self, f, *partial) -> None: ...
+
+    @method(egg_fn="unstable-app")
+    def __call__(self, *args: Unpack[TS]) -> T: ...
+
+
+converter(RuntimeFunction, UnstableFn, UnstableFn)
+converter(partial, UnstableFn, lambda p: UnstableFn(p.func, *p.args))

egglog/conversion.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, NewType, TypeVar, cast
+
+from .declarations import *
+from .pretty import *
+from .runtime import *
+from .thunk import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from .declarations import HasDeclerations
+    from .egraph import Expr
+
+__all__ = ["convert", "converter", "resolve_literal", "convert_to_same_type"]
+# Mapping from (source type, target type) to and function which takes in the runtimes values of the source and return the target
+TypeName = NewType("TypeName", str)
+CONVERSIONS: dict[tuple[type | TypeName, TypeName], tuple[int, Callable]] = {}
+# Global declerations to store all convertable types so we can query if they have certain methods or not
+# Defer it as a thunk so we can register conversions without triggering type signature loading
+CONVERSIONS_DECLS: Callable[[], Declarations] = Thunk.value(Declarations())
+
+T = TypeVar("T")
+V = TypeVar("V", bound="Expr")
+
+
+class ConvertError(Exception):
+    pass
+
+
+def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost: int = 1) -> None:
+    """
+    Register a converter from some type to an egglog type.
+    """
+    to_type_name = process_tp(to_type)
+    if not isinstance(to_type_name, str):
+        raise TypeError(f"Expected return type to be a egglog type, got {to_type_name}")
+    _register_converter(process_tp(from_type), to_type_name, fn, cost)
+
+
+def _register_converter(a: type | TypeName, b: TypeName, a_b: Callable, cost: int) -> None:
+    """
+    Registers a converter from some type to an egglog type, if not already registered.
+
+    Also adds transitive converters, i.e. if registering A->B and there is already B->C, then A->C will be registered.
+    Also, if registering A->B and there is already D->A, then D->B will be registered.
+    """
+    if a == b:
+        return
+    if (a, b) in CONVERSIONS and CONVERSIONS[(a, b)][0] <= cost:
+        return
+    CONVERSIONS[(a, b)] = (cost, a_b)
+    for (c, d), (other_cost, c_d) in list(CONVERSIONS.items()):
+        if b == c:
+            _register_converter(a, d, _ComposedConverter(a_b, c_d), cost + other_cost)
+        if a == d:
+            _register_converter(c, b, _ComposedConverter(c_d, a_b), cost + other_cost)
+
+
+@dataclass
+class _ComposedConverter:
+    """
+    A converter which is composed of multiple converters.
+
+    _ComposeConverter(a_b, b_c) is equivalent to lambda x: b_c(a_b(x))
+
+    We use the dataclass instead of the lambda to make it easier to debug.
+    """
+
+    a_b: Callable
+    b_c: Callable
+
+    def __call__(self, x: object) -> object:
+        return self.b_c(self.a_b(x))
+
+    def __str__(self) -> str:
+        return f"{self.b_c} ∘ {self.a_b}"
+
+
+def convert(source: object, target: type[V]) -> V:
+    """
+    Convert a source object to a target type.
+    """
+    assert isinstance(target, RuntimeClass)
+    return cast(V, resolve_literal(target.__egg_tp__, source))
+
+
+def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr:
+    """
+    Convert a source object to the same type as the target.
+    """
+    tp = target.__egg_typed_expr__.tp
+    return resolve_literal(tp.to_var(), source)
+
+
+def process_tp(tp: type | RuntimeClass) -> TypeName | type:
+    """
+    Process a type before converting it, to add it to the global declerations and resolve to a ref.
+    """
+    global CONVERSIONS_DECLS
+    if isinstance(tp, RuntimeClass):
+        CONVERSIONS_DECLS = Thunk.fn(_combine_decls, CONVERSIONS_DECLS, tp)
+        egg_tp = tp.__egg_tp__
+        if egg_tp.args:
+            raise TypeError(f"Cannot register a converter for a generic type, got {tp}")
+        return TypeName(egg_tp.name)
+    return tp
+
+
+def _combine_decls(d: Callable[[], Declarations], x: HasDeclerations) -> Declarations:
+    return Declarations.create(d(), x)
+
+
+def min_convertable_tp(a: object, b: object, name: str) -> TypeName:
+    """
+    Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
+    """
+    decls = CONVERSIONS_DECLS()
+    a_tp = _get_tp(a)
+    b_tp = _get_tp(b)
+    a_converts_to = {
+        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to, name)
+    }
+    b_converts_to = {
+        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to, name)
+    }
+    if isinstance(a_tp, str):
+        a_converts_to[a_tp] = 0
+    if isinstance(b_tp, str):
+        b_converts_to[b_tp] = 0
+    common = set(a_converts_to) & set(b_converts_to)
+    if not common:
+        raise ConvertError(f"Cannot convert {a_tp} and {b_tp} to a common type")
+    return min(common, key=lambda tp: a_converts_to[tp] + b_converts_to[tp])
+
+
+def identity(x: object) -> object:
+    return x
+
+
+def resolve_literal(tp: TypeOrVarRef, arg: object) -> RuntimeExpr:
+    arg_type = _get_tp(arg)
+
+    # If we have any type variables, dont bother trying to resolve the literal, just return the arg
+    try:
+        tp_just = tp.to_just()
+    except NotImplementedError:
+        # If this is a var, it has to be a runtime expession
+        assert isinstance(arg, RuntimeExpr)
+        return arg
+    tp_name = TypeName(tp_just.name)
+    if arg_type == tp_name:
+        # If the type is an egg type, it has to be a runtime expr
+        assert isinstance(arg, RuntimeExpr)
+        return arg
+    # Try all parent types as well, if we are converting from a Python type
+    for arg_type_instance in arg_type.__mro__ if isinstance(arg_type, type) else [arg_type]:
+        try:
+            fn = CONVERSIONS[(cast(TypeName | type, arg_type_instance), tp_name)][1]
+        except KeyError:
+            continue
+        break
+    else:
+        raise ConvertError(f"Cannot convert {arg_type} to {tp_name}")
+    return fn(arg)
+
+
+def _get_tp(x: object) -> TypeName | type:
+    if isinstance(x, RuntimeExpr):
+        return TypeName(x.__egg_typed_expr__.tp.name)
+    tp = type(x)
+    # If this value has a custom metaclass, let's use that as our index instead of the type
+    if type(tp) != type:
+        return type(tp)
+    return tp