egglog 10.0.2__cp312-cp312-win_amd64.whl → 11.1.0__cp312-cp312-win_amd64.whl

egglog/conversion.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 from .declarations import *
 from .pretty import *
@@ -13,14 +14,14 @@ from .thunk import *
 from .type_constraint_solver import TypeConstraintError
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Generator
+    from collections.abc import Generator
 
     from .egraph import BaseExpr
     from .type_constraint_solver import TypeConstraintSolver
 
-__all__ = ["ConvertError", "convert", "convert_to_same_type", "converter", "resolve_literal"]
+__all__ = ["ConvertError", "convert", "converter", "get_type_args"]
 # Mapping from (source type, target type) to and function which takes in the runtimes values of the source and return the target
-CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable]] = {}
+CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable[[Any], RuntimeExpr]]] = {}
 # Global declerations to store all convertable types so we can query if they have certain methods or not
 _CONVERSION_DECLS = Declarations.create()
 # Defer a list of declerations to be added to the global declerations, so that we can not trigger them procesing
@@ -28,7 +29,7 @@ _CONVERSION_DECLS = Declarations.create()
 _TO_PROCESS_DECLS: list[DeclerationsLike] = []
 
 
-def _retrieve_conversion_decls() -> Declarations:
+def retrieve_conversion_decls() -> Declarations:
     _CONVERSION_DECLS.update(*_TO_PROCESS_DECLS)
     _TO_PROCESS_DECLS.clear()
     return _CONVERSION_DECLS
@@ -49,10 +50,10 @@ def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost:
     to_type_name = process_tp(to_type)
     if not isinstance(to_type_name, JustTypeRef):
         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)
+    _register_converter(process_tp(from_type), to_type_name, cast("Callable[[Any], RuntimeExpr]", fn), cost)
 
 
-def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable, cost: int) -> None:
+def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable[[Any], RuntimeExpr], cost: int) -> None:
     """
     Registers a converter from some type to an egglog type, if not already registered.
 
@@ -97,15 +98,15 @@ class _ComposedConverter:
     We use the dataclass instead of the lambda to make it easier to debug.
     """
 
-    a_b: Callable
-    b_c: Callable
+    a_b: Callable[[Any], RuntimeExpr]
+    b_c: Callable[[Any], RuntimeExpr]
     b_args: tuple[JustTypeRef, ...]
 
-    def __call__(self, x: object) -> object:
+    def __call__(self, x: Any) -> RuntimeExpr:
         # if we have A -> B and B[C] -> D then we should use (C,) as the type args
         # when converting from A -> B
         if self.b_args:
-            with with_type_args(self.b_args, _retrieve_conversion_decls):
+            with with_type_args(self.b_args, retrieve_conversion_decls):
                 first_res = self.a_b(x)
         else:
             first_res = self.a_b(x)
@@ -142,36 +143,53 @@ def process_tp(tp: type | RuntimeClass) -> JustTypeRef | type:
     return tp
 
 
-def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
+def min_binary_conversion(
+    method_name: str, lhs: type | JustTypeRef, rhs: type | JustTypeRef
+) -> tuple[Callable[[Any], RuntimeExpr], Callable[[Any], RuntimeExpr]] | None:
     """
-    Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
+    Given a binary method and two starting types for the LHS and RHS, return a pair of callable which will convert
+    the LHS and RHS to appropriate types which support this method. If no such conversion is possible, return None.
+
+    It should return the types which minimize the total conversion cost. If one of the types is a Python type, then
+    both of them can be converted. However, if both are egglog types, then only one of them can be converted.
+    """
+    decls = retrieve_conversion_decls()
+    # tuple of (cost, convert lhs, convert rhs)
+    best_method: tuple[int, Callable[[Any], RuntimeExpr], Callable[[Any], RuntimeExpr]] | None = None
+
+    possible_lhs = _all_conversions_from(lhs) if isinstance(lhs, type) else [(0, lhs, identity)]
+    possible_rhs = _all_conversions_from(rhs) if isinstance(rhs, type) else [(0, rhs, identity)]
+    for lhs_cost, lhs_converted_type, lhs_convert in possible_lhs:
+        # Start by checking if we have a LHS that matches exactly and a RHS which can be converted
+        if (desired_other_type := decls.check_binary_method_with_self_type(method_name, lhs_converted_type)) and (
+            converter := CONVERSIONS.get((rhs, desired_other_type))
+        ):
+            cost = lhs_cost + converter[0]
+            if best_method is None or best_method[0] > cost:
+                best_method = (cost, lhs_convert, converter[1])
+
+    for rhs_cost, rhs_converted_type, rhs_convert in possible_rhs:
+        # Next see if it's possible to convert the LHS and keep the RHS as is
+        for desired_self_type in decls.check_binary_method_with_other_type(method_name, rhs_converted_type):
+            if converter := CONVERSIONS.get((lhs, desired_self_type)):
+                cost = rhs_cost + converter[0]
+                if best_method is None or best_method[0] > cost:
+                    best_method = (cost, converter[1], rhs_convert)
+    if best_method is None:
+        return None
+    return best_method[1], best_method[2]
+
+
+def _all_conversions_from(tp: JustTypeRef | type) -> list[tuple[int, JustTypeRef, Callable[[Any], RuntimeExpr]]]:
+    """
+    Get all conversions from a type to other types.
+
+    Returns a list of tuples of (cost, target type, conversion function).
     """
-    decls = _retrieve_conversion_decls()
-    a_tp = _get_tp(a)
-    b_tp = _get_tp(b)
-    # Make sure at least one of the types has the method, to avoid issue with == upcasting improperly
-    if not (
-        (isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name))
-        or (isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name))
-    ):
-        raise ConvertError(f"Neither {a_tp} nor {b_tp} has method {name}")
-    a_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == a_tp and decls.has_method(to.name, name)
-    }
-    b_converts_to = {
-        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to.name, name)
-    }
-    if isinstance(a_tp, JustTypeRef) and decls.has_method(a_tp.name, name):
-        a_converts_to[a_tp] = 0
-    if isinstance(b_tp, JustTypeRef) and decls.has_method(b_tp.name, name):
-        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 [(cost, target, fn) for (source, target), (cost, fn) in CONVERSIONS.items() if source == tp]
+
+
+def identity(x: Any) -> Any:
     return x
 
 
@@ -197,7 +215,7 @@ def with_type_args(args: tuple[JustTypeRef, ...], decls: Callable[[], Declaratio
 def resolve_literal(
     tp: TypeOrVarRef,
     arg: object,
-    decls: Callable[[], Declarations] = _retrieve_conversion_decls,
+    decls: Callable[[], Declarations] = retrieve_conversion_decls,
     tcs: TypeConstraintSolver | None = None,
     cls_name: str | None = None,
 ) -> RuntimeExpr:
@@ -208,12 +226,12 @@ def resolve_literal(
 
     If it cannot be resolved, we assume that the value passed in will resolve it.
     """
-    arg_type = _get_tp(arg)
+    arg_type = resolve_type(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:
+    except TypeVarError:
         # If this is a generic arg but passed in a non runtime expression, try to resolve the generic
         # args first based on the existing type constraint solver
         if tcs:
@@ -258,7 +276,7 @@ def _debug_print_converers():
         source_to_targets[source].append(target)
 
 
-def _get_tp(x: object) -> JustTypeRef | type:
+def resolve_type(x: object) -> JustTypeRef | type:
     if isinstance(x, RuntimeExpr):
         return x.__egg_typed_expr__.tp
     tp = type(x)

egglog/declarations.py

@@ -51,6 +51,7 @@ __all__ = [
     "InitRef",
     "JustTypeRef",
     "LetDecl",
+    "LetRefDecl",
     "LitDecl",
     "LitType",
     "MethodRef",
@@ -72,16 +73,18 @@ __all__ = [
     "SpecialFunctions",
     "TypeOrVarRef",
     "TypeRefWithVars",
+    "TypeVarError",
     "TypedExprDecl",
+    "UnboundVarDecl",
     "UnionDecl",
     "UnnamedFunctionRef",
-    "VarDecl",
+    "collect_unbound_vars",
     "replace_typed_expr",
     "upcast_declerations",
 ]
 
 
-@dataclass
+@dataclass(match_args=False)
 class DelayedDeclerations:
     __egg_decls_thunk__: Callable[[], Declarations] = field(repr=False)
 
@@ -93,7 +96,7 @@ class DelayedDeclerations:
         # Catch attribute error, so that it isn't bubbled up as a missing attribute and fallbacks on `__getattr__`
         # instead raise explicitly
         except AttributeError as err:
-            msg = f"Cannot resolve declarations for {self}"
+            msg = f"Cannot resolve declarations for {self}: {err}"
             raise RuntimeError(msg) from err
 
 
@@ -223,14 +226,46 @@ class Declarations:
             case _:
                 assert_never(ref)
 
-    def has_method(self, class_name: str, method_name: str) -> bool | None:
+    def check_binary_method_with_types(self, method_name: str, self_type: JustTypeRef, other_type: JustTypeRef) -> bool:
+        """
+        Checks if the class has a binary method compatible with the given types.
+        """
+        vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+        if callable_decl := self._classes[self_type.name].methods.get(method_name):
+            match callable_decl.signature:
+                case FunctionSignature((self_arg_type, other_arg_type)) if self_arg_type.matches_just(
+                    vars, self_type
+                ) and other_arg_type.matches_just(vars, other_type):
+                    return True
+        return False
+
+    def check_binary_method_with_self_type(self, method_name: str, self_type: JustTypeRef) -> JustTypeRef | None:
         """
-        Returns whether the given class has the given method, or None if we cant find the class.
+        Checks if the class has a binary method with the given name and self type. Returns the other type if it exists.
         """
-        if class_name in self._classes:
-            return method_name in self._classes[class_name].methods
+        vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+        class_decl = self._classes.get(self_type.name)
+        if class_decl is None:
+            return None
+        if callable_decl := class_decl.methods.get(method_name):
+            match callable_decl.signature:
+                case FunctionSignature((self_arg_type, other_arg_type)) if self_arg_type.matches_just(vars, self_type):
+                    return other_arg_type.to_just(vars)
         return None
 
+    def check_binary_method_with_other_type(self, method_name: str, other_type: JustTypeRef) -> Iterable[JustTypeRef]:
+        """
+        Returns the types which are compatible with the given binary method name and other type.
+        """
+        for class_decl in self._classes.values():
+            vars: dict[ClassTypeVarRef, JustTypeRef] = {}
+            if callable_decl := class_decl.methods.get(method_name):
+                match callable_decl.signature:
+                    case FunctionSignature((self_arg_type, other_arg_type)) if other_arg_type.matches_just(
+                        vars, other_type
+                    ):
+                        yield self_arg_type.to_just(vars)
+
     def get_class_decl(self, name: str) -> ClassDecl:
         return self._classes[name]
 
@@ -254,6 +289,7 @@ class ClassDecl:
     methods: dict[str, FunctionDecl | ConstructorDecl] = field(default_factory=dict)
     properties: dict[str, FunctionDecl | ConstructorDecl] = field(default_factory=dict)
     preserved_methods: dict[str, Callable] = field(default_factory=dict)
+    match_args: tuple[str, ...] = field(default=())
 
 
 @dataclass(frozen=True)
@@ -298,6 +334,10 @@ class JustTypeRef:
 _RESOLVED_TYPEVARS: dict[ClassTypeVarRef, TypeVar] = {}
 
 
+class TypeVarError(RuntimeError):
+    """Error when trying to resolve a type variable that doesn't exist."""
+
+
 @dataclass(frozen=True)
 class ClassTypeVarRef:
     """
@@ -307,9 +347,10 @@ class ClassTypeVarRef:
     name: str
     module: str
 
-    def to_just(self) -> JustTypeRef:
-        msg = f"{self}: egglog does not support generic classes yet."
-        raise NotImplementedError(msg)
+    def to_just(self, vars: dict[ClassTypeVarRef, JustTypeRef] | None = None) -> JustTypeRef:
+        if vars is None or self not in vars:
+            raise TypeVarError(f"Cannot convert type variable {self} to concrete type without variable bindings")
+        return vars[self]
 
     def __str__(self) -> str:
         return str(self.to_type_var())
@@ -323,20 +364,39 @@ class ClassTypeVarRef:
     def to_type_var(self) -> TypeVar:
         return _RESOLVED_TYPEVARS[self]
 
+    def matches_just(self, vars: dict[ClassTypeVarRef, JustTypeRef], other: JustTypeRef) -> bool:
+        """
+        Checks if this type variable matches the given JustTypeRef, including type variables.
+        """
+        if self in vars:
+            return vars[self] == other
+        vars[self] = other
+        return True
+
 
 @dataclass(frozen=True)
 class TypeRefWithVars:
     name: str
     args: tuple[TypeOrVarRef, ...] = ()
 
-    def to_just(self) -> JustTypeRef:
-        return JustTypeRef(self.name, tuple(a.to_just() for a in self.args))
+    def to_just(self, vars: dict[ClassTypeVarRef, JustTypeRef] | None = None) -> JustTypeRef:
+        return JustTypeRef(self.name, tuple(a.to_just(vars) for a in self.args))
 
     def __str__(self) -> str:
         if self.args:
             return f"{self.name}[{', '.join(str(a) for a in self.args)}]"
         return self.name
 
+    def matches_just(self, vars: dict[ClassTypeVarRef, JustTypeRef], other: JustTypeRef) -> bool:
+        """
+        Checks if this type reference matches the given JustTypeRef, including type variables.
+        """
+        return (
+            self.name == other.name
+            and len(self.args) == len(other.args)
+            and all(a.matches_just(vars, b) for a, b in zip(self.args, other.args, strict=True))
+        )
+
 
 TypeOrVarRef: TypeAlias = ClassTypeVarRef | TypeRefWithVars
 
@@ -361,7 +421,7 @@ class UnnamedFunctionRef:
         arg_names = []
         for a in self.args:
             arg_types.append(a.tp.to_var())
-            assert isinstance(a.expr, VarDecl)
+            assert isinstance(a.expr, UnboundVarDecl)
             arg_names.append(a.expr.name)
         return FunctionSignature(
             arg_types=tuple(arg_types),
@@ -514,10 +574,14 @@ CallableDecl: TypeAlias = RelationDecl | ConstantDecl | FunctionDecl | Construct
 
 
 @dataclass(frozen=True)
-class VarDecl:
+class UnboundVarDecl:
+    name: str
+    egg_name: str | None = None
+
+
+@dataclass(frozen=True)
+class LetRefDecl:
     name: str
-    # Differentiate between let bound vars and vars created in rules so that they won't shadow in egglog, by adding a prefix
-    is_let: bool
 
 
 @dataclass(frozen=True)
@@ -628,7 +692,7 @@ class PartialCallDecl:
     call: CallDecl
 
 
-ExprDecl: TypeAlias = VarDecl | LitDecl | CallDecl | PyObjectDecl | PartialCallDecl
+ExprDecl: TypeAlias = UnboundVarDecl | LetRefDecl | LitDecl | CallDecl | PyObjectDecl | PartialCallDecl
 
 
 @dataclass(frozen=True)
@@ -678,6 +742,28 @@ def replace_typed_expr(typed_expr: TypedExprDecl, replacements: Mapping[TypedExp
     return _inner(typed_expr)
 
 
+def collect_unbound_vars(typed_expr: TypedExprDecl) -> set[TypedExprDecl]:
+    """
+    Returns the set of all unbound vars
+    """
+    seen = set[TypedExprDecl]()
+    unbound_vars = set[TypedExprDecl]()
+
+    def visit(typed_expr: TypedExprDecl) -> None:
+        if typed_expr in seen:
+            return
+        seen.add(typed_expr)
+        match typed_expr.expr:
+            case CallDecl(_, args) | PartialCallDecl(CallDecl(_, args)):
+                for arg in args:
+                    visit(arg)
+            case UnboundVarDecl(_):
+                unbound_vars.add(typed_expr)
+
+    visit(typed_expr)
+    return unbound_vars
+
+
 ##
 # Schedules
 ##

egglog/deconstruct.py

@@ -0,0 +1,173 @@
+"""
+Utility functions to deconstruct expressions in Python.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from functools import partial
+from typing import TYPE_CHECKING, TypeVar, overload
+
+from typing_extensions import TypeVarTuple, Unpack
+
+from .declarations import *
+from .egraph import BaseExpr
+from .runtime import *
+from .thunk import *
+
+if TYPE_CHECKING:
+    from .builtins import Bool, PyObject, String, UnstableFn, f64, i64
+
+
+T = TypeVar("T", bound=BaseExpr)
+TS = TypeVarTuple("TS", default=Unpack[tuple[BaseExpr, ...]])
+
+__all__ = ["get_callable_args", "get_callable_fn", "get_let_name", "get_literal_value", "get_var_name"]
+
+
+@overload
+def get_literal_value(x: String) -> str | None: ...
+
+
+@overload
+def get_literal_value(x: Bool) -> bool | None: ...
+
+
+@overload
+def get_literal_value(x: i64) -> int | None: ...
+
+
+@overload
+def get_literal_value(x: f64) -> float | None: ...
+
+
+@overload
+def get_literal_value(x: PyObject) -> object: ...
+
+
+@overload
+def get_literal_value(x: UnstableFn[T, Unpack[TS]]) -> Callable[[Unpack[TS]], T] | None: ...
+
+
+def get_literal_value(x: String | Bool | i64 | f64 | PyObject | UnstableFn) -> object:
+    """
+    Returns the literal value of an expression if it is a literal.
+    If it is not a literal, returns None.
+    """
+    if not isinstance(x, RuntimeExpr):
+        raise TypeError(f"Expected Expression, got {type(x).__name__}")
+    match x.__egg_typed_expr__.expr:
+        case LitDecl(v):
+            return v
+        case PyObjectDecl(obj):
+            return obj
+        case PartialCallDecl(call):
+            fn, args = _deconstruct_call_decl(x.__egg_decls_thunk__, call)
+            if not args:
+                return fn
+            return partial(fn, *args)
+    return None
+
+
+def get_let_name(x: BaseExpr) -> str | None:
+    """
+    Check if the expression is a `let` expression and return the name of the variable.
+    If it is not a `let` expression, return None.
+    """
+    if not isinstance(x, RuntimeExpr):
+        raise TypeError(f"Expected Expression, got {type(x).__name__}")
+    match x.__egg_typed_expr__.expr:
+        case LetRefDecl(name):
+            return name
+    return None
+
+
+def get_var_name(x: BaseExpr) -> str | None:
+    """
+    Check if the expression is a variable and return its name.
+    If it is not a variable, return None.
+    """
+    if not isinstance(x, RuntimeExpr):
+        raise TypeError(f"Expected Expression, got {type(x).__name__}")
+    match x.__egg_typed_expr__.expr:
+        case UnboundVarDecl(name, _egg_name):
+            return name
+    return None
+
+
+def get_callable_fn(x: T) -> Callable[..., T] | None:
+    """
+    Gets the function of an expression if it is a call expression.
+    If it is not a call expression (a property, a primitive value, constants, classvars, a let value), return None.
+    For those values, you can check them by comparing them directly with equality or for primitives calling `.eval()`
+    to return the Python value.
+    """
+    if not isinstance(x, RuntimeExpr):
+        raise TypeError(f"Expected Expression, got {type(x).__name__}")
+    match x.__egg_typed_expr__.expr:
+        case CallDecl() as call:
+            fn, _ = _deconstruct_call_decl(x.__egg_decls_thunk__, call)
+            return fn
+    return None
+
+
+@overload
+def get_callable_args(x: T, fn: None = ...) -> tuple[BaseExpr, ...]: ...
+
+
+@overload
+def get_callable_args(x: T, fn: Callable[[Unpack[TS]], T]) -> tuple[Unpack[TS]] | None: ...
+
+
+def get_callable_args(x: T, fn: Callable[[Unpack[TS]], T] | None = None) -> tuple[Unpack[TS]] | None:
+    """
+    Gets all the arguments of an expression.
+    If a function is provided, it will only return the arguments if the expression is a call
+    to that function.
+
+    Note that recursively calling the arguments is the safe way to walk the expression tree.
+    """
+    if not isinstance(x, RuntimeExpr):
+        raise TypeError(f"Expected Expression, got {type(x).__name__}")
+    match x.__egg_typed_expr__.expr:
+        case CallDecl() as call:
+            actual_fn, args = _deconstruct_call_decl(x.__egg_decls_thunk__, call)
+            if fn is None:
+                return args
+            # Compare functions and classes without considering bound type parameters, so that you can pass
+            # in a binding like Vec[i64] and match Vec[i64](...) or Vec(...) calls.
+            if (
+                isinstance(actual_fn, RuntimeFunction)
+                and isinstance(fn, RuntimeFunction)
+                and actual_fn.__egg_ref__ == fn.__egg_ref__
+            ):
+                return args
+            if (
+                isinstance(actual_fn, RuntimeClass)
+                and isinstance(fn, RuntimeClass)
+                and actual_fn.__egg_tp__.name == fn.__egg_tp__.name
+            ):
+                return args
+    return None
+
+
+def _deconstruct_call_decl(
+    decls_thunk: Callable[[], Declarations], call: CallDecl
+) -> tuple[Callable, tuple[object, ...]]:
+    """
+    Deconstructs a CallDecl into a runtime callable and its arguments.
+    """
+    args = call.args
+    arg_exprs = tuple(RuntimeExpr(decls_thunk, Thunk.value(a)) for a in args)
+    if isinstance(call.callable, InitRef):
+        return RuntimeClass(
+            decls_thunk,
+            TypeRefWithVars(call.callable.class_name, tuple(tp.to_var() for tp in (call.bound_tp_params or []))),
+        ), arg_exprs
+    egg_bound = (
+        JustTypeRef(call.callable.class_name, call.bound_tp_params or ())
+        if isinstance(call.callable, ClassMethodRef)
+        else None
+    )
+
+    return RuntimeFunction(decls_thunk, Thunk.value(call.callable), egg_bound), arg_exprs