Try fixing type resolution again

Author: saulshanabrook

python/egglog/conversion.py

@@ -11,13 +11,11 @@
 from .pretty import *
 from .runtime import *
 from .thunk import *
-from .type_constraint_solver import TypeConstraintError
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
     from .egraph import BaseExpr
-    from .type_constraint_solver import TypeConstraintSolver
 
 __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
@@ -220,41 +218,17 @@ def resolve_literal(
     tp: TypeOrVarRef,
     arg: object,
     decls: Callable[[], Declarations] = retrieve_conversion_decls,
-    tcs: TypeConstraintSolver | None = None,
 ) -> RuntimeExpr:
     """
     Try to convert an object to a type, raising a ConvertError if it is not possible.
 
-    If the type has vars in it, they will be tried to be resolved into concrete vars based on the type constraint solver.
-
     If it cannot be resolved, we assume that the value passed in will resolve it.
     """
-    arg_type = resolve_type(arg)
-
-    # If we have any type variables, don't bother trying to resolve the literal, just return the arg
-    try:
-        tp_just = tp.to_just()
-    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:
-            try:
-                tp_just = tcs.substitute_typevars_try_function(tp, arg, decls)
-            # If we can't resolve the type var yet, then just assume it is the right value
-            except TypeConstraintError as e:
-                if not isinstance(arg, RuntimeExpr):
-                    raise ConvertError(f"Cannot convert {arg} of type {arg_type} to {tp}") from e
-                tp_just = arg.__egg_typed_expr__.tp
-        else:
-            # If this is a var, it has to be a runtime expression
-            assert isinstance(arg, RuntimeExpr), f"Expected a runtime expression, got {arg}"
-            return arg
-    if tcs:
-        tcs.infer_typevars(tp, tp_just)
-    if arg_type == tp_just:
-        # If the type is an egg type, it has to be a runtime expr
-        assert isinstance(arg, RuntimeExpr)
+    # If this is a runtime expression that could match the type already, just return it
+    if isinstance(arg, RuntimeExpr) and tp.matches_just({}, arg.__egg_typed_expr__.tp):
         return arg
+    tp_just = tp.to_just()
+    arg_type = resolve_type(arg)
     if arg is DUMMY_VALUE:
         return RuntimeExpr.__from_values__(decls(), TypedExprDecl(tp_just, DummyDecl()))
     if (conversion := _lookup_conversion(arg_type, tp_just)) is not None:

python/egglog/declarations.py

@@ -8,6 +8,7 @@
 
 from dataclasses import dataclass, field
 from functools import cached_property
+from itertools import chain, repeat
 from typing import (
     TYPE_CHECKING,
     ClassVar,
@@ -413,6 +414,13 @@ def matches_just(self, vars: dict[TypeVarRef, JustTypeRef], other: JustTypeRef)
         vars[self] = other
         return True
 
+    @property
+    def vars(self) -> set[TypeVarRef]:
+        """
+        Returns all type variables in this type reference.
+        """
+        return {self}
+
 
 @dataclass(frozen=True)
 class TypeRefWithVars:
@@ -437,6 +445,16 @@ def matches_just(self, vars: dict[TypeVarRef, JustTypeRef], other: JustTypeRef)
             and all(a.matches_just(vars, b) for a, b in zip(self.args, other.args, strict=True))
         )
 
+    @property
+    def vars(self) -> set[TypeVarRef]:
+        """
+        Returns all type variables in this type reference.
+        """
+        vars = set[TypeVarRef]()
+        for arg in self.args:
+            vars.update(arg.vars)
+        return vars
+
 
 TypeOrVarRef: TypeAlias = TypeVarRef | TypeRefWithVars
 
@@ -589,6 +607,25 @@ def semantic_return_type(self) -> TypeOrVarRef:
     def mutates(self) -> bool:
         return self.return_type is None
 
+    @property
+    def arg_vars(self) -> set[TypeVarRef]:
+        """
+        Returns all type variables in the argument types.
+        """
+        vars = set[TypeVarRef]()
+        for arg in self.arg_types:
+            vars.update(arg.vars)
+        if self.var_arg_type:
+            vars.update(self.var_arg_type.vars)
+        return vars
+
+    @property
+    def all_args(self) -> Iterable[TypeOrVarRef]:
+        """
+        Returns all argument types, including var args.
+        """
+        return chain(self.arg_types, (repeat(self.var_arg_type) if self.var_arg_type else []))
+
 
 @dataclass(frozen=True)
 class FunctionDecl:
@@ -675,7 +712,7 @@ def __new__(cls, *args: object, **kwargs: object) -> Self:
         """
         Pool CallDecls so that they can be compared by identity more quickly.
 
-        Neccessary bc we search for common parents when serializing CallDecl trees to egglog to
+        Necessary bc we search for common parents when serializing CallDecl trees to egglog to
         only serialize each sub-tree once.
         """
         # normalize the args/kwargs to a tuple so that they can be compared

python/egglog/deconstruct.py

@@ -188,8 +188,8 @@ def _deconstruct_call_decl(
             TypeRefWithVars(call.callable.ident, tuple(tp.to_var() for tp in (call.bound_tp_params or []))),
         ), arg_exprs
     egg_bound = (
-        JustTypeRef(call.callable.ident, call.bound_tp_params or ())
-        if isinstance(call.callable, (ClassMethodRef, MethodRef))
+        JustTypeRef(call.callable.ident, call.bound_tp_params)
+        if isinstance(call.callable, (ClassMethodRef, MethodRef, InitRef)) and call.bound_tp_params
         else None
     )
 

python/egglog/egraph_state.py

@@ -461,7 +461,9 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str:
         decl = self.__egg_decls__._classes[ref.ident]
         self.type_ref_to_egg_sort[ref] = egg_name = (not ref.args and 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 decl.builtin:
+            # If this has args, create a new parameterized version of the builtin class
             if ref.args:
                 if ref.ident == Ident.builtin("UnstableFn"):
                     # UnstableFn is a special case, where the rest of args are collected into a call
@@ -478,18 +480,18 @@ def type_ref_to_egg(self, ref: JustTypeRef) -> str:
                     ]
                 else:
                     type_args = [bindings.Var(span(), self.type_ref_to_egg(a)) for a in ref.args]
-                args = (self.type_ref_to_egg(JustTypeRef(ref.ident)), type_args)
-            else:
-                args = None
-            self.egraph.run_program(bindings.Sort(span(), egg_name, args))
-        # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods, because
-        # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted
-        # even if you never use that function.
-        if decl.builtin:
+                assert decl.egg_name
+                self.egraph.run_program(bindings.Sort(span(), egg_name, (decl.egg_name, type_args)))
+
+            # For builtin classes, let's also make sure we have the mapping of all egg fn names for class methods.
+            # these can be created even without adding them to the e-graph, like `vec-empty` which can be extracted
+            # even if you never use that function.
             for method_name in decl.class_methods:
                 self.callable_ref_to_egg(ClassMethodRef(ref.ident, method_name))
             if decl.init:
                 self.callable_ref_to_egg(InitRef(ref.ident))
+        else:
+            self.egraph.run_program(bindings.Sort(span(), egg_name, None))
 
         return egg_name
 
@@ -760,23 +762,23 @@ def value_to_expr(self, tp: JustTypeRef, value: bindings.Value) -> ExprDecl:  #
                 return CallDecl(
                     InitRef(Ident.builtin("Set")),
                     tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs_),
-                    (v_tp,),
+                    (v_tp,) if not xs_ else (),
                 )
             case "Vec":
                 xs = self.egraph.value_to_vec(value)
                 (v_tp,) = tp.args
                 return CallDecl(
                     InitRef(Ident.builtin("Vec")),
                     tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs),
-                    (v_tp,),
+                    (v_tp,) if not xs else (),
                 )
             case "MultiSet":
                 xs = self.egraph.value_to_multiset(value)
                 (v_tp,) = tp.args
                 return CallDecl(
                     InitRef(Ident.builtin("MultiSet")),
                     tuple(TypedExprDecl(v_tp, self.value_to_expr(v_tp, x)) for x in xs),
-                    (v_tp,),
+                    (v_tp,) if not xs else (),
                 )
             case "UnstableFn":
                 _names, _args = self.egraph.value_to_function(value)
@@ -796,22 +798,13 @@ def _unstable_fn_value_to_expr(
                 continue
             if signature.semantic_return_type.ident != return_tp.ident:
                 continue
-            tcs = TypeConstraintSolver()
-
-            arg_types = tcs.infer_arg_types(
+            arg_types = TypeConstraintSolver().infer_arg_types(
                 signature.arg_types, signature.semantic_return_type, signature.var_arg_type, return_tp
             )
-
             args = tuple(
                 TypedExprDecl(tp, self.value_to_expr(tp, v)) for tp, v in zip(arg_types, partial_args, strict=False)
             )
-            if isinstance(callable_ref, ClassMethodRef | InitRef):
-                bound_tp_params = tuple(
-                    map(tcs.substitute_typevars, self.__egg_decls__.get_class_decl(callable_ref.ident).type_vars)
-                )
-            else:
-                bound_tp_params = ()
-            call_decl = CallDecl(callable_ref, args, bound_tp_params)
+            call_decl = CallDecl(callable_ref, args)
             return PartialCallDecl(call_decl)
         raise ValueError(f"Function '{name}' not found")
 
@@ -936,6 +929,9 @@ def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl:
                 tcs = TypeConstraintSolver()
                 if possible_type and possible_type.args:
                     tcs.bind_class(possible_type, self.decls)
+                    bound_args = possible_type.args
+                else:
+                    bound_args = ()
                 try:
                     arg_types = tcs.infer_arg_types(
                         signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp
@@ -945,12 +941,9 @@ def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl:
                 except TypeConstraintError:
                     continue
                 args = tuple(self.resolve_term(a, tp) for a, tp in a_tp)
-                if not args and isinstance(callable_ref, ClassMethodRef | InitRef):
-                    bound_tp_params = tuple(
-                        map(tcs.substitute_typevars, self.decls.get_class_decl(callable_ref.ident).type_vars)
-                    )
-                else:
-                    bound_tp_params = ()
+                # Only save bound tp params if needed for inferring return type
+                # this is true if the set of set of type vars in the return are not a subset of those in the args
+                bound_tp_params = () if signature.semantic_return_type.vars.issubset(signature.arg_vars) else bound_args
                 return CallDecl(callable_ref, args, bound_tp_params)
         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]}"

python/egglog/exp/array_api.py

@@ -2500,7 +2500,7 @@ def factor_ruleset(
     yield rule(
         eq(n).to(polynomial(mss)),
         # Find factor that shows up in most monomials, at least two of them
-        counts == MultiSet.sum_multisets(mss.map(MultiSet[Value].reset_counts)),
+        counts == MultiSet.sum_multisets(mss.map(MultiSet.reset_counts)),
         eq(factor).to(counts.pick_max()),
         # Only factor out if it term appears in more than one monomial
         counts.count(factor) > 1,

python/egglog/exp/polynomials.py

@@ -253,7 +253,7 @@ def factor_ruleset(
     yield rule(
         n == polynomial(mss),
         # Find factor that shows up in most monomials, at least two of them
-        counts == MultiSet.sum_multisets(mss.map(MultiSet[Number].reset_counts)),
+        counts == MultiSet.sum_multisets(mss.map(MultiSet.reset_counts)),
         factor == counts.pick_max(),
         # Only factor out if it term appears in more than one monomial
         counts.count(factor) > 1,

python/egglog/runtime.py

@@ -17,7 +17,6 @@
 from collections.abc import Callable
 from dataclasses import InitVar, dataclass, replace
 from inspect import Parameter, Signature
-from itertools import zip_longest
 from typing import TYPE_CHECKING, Any, TypeVar, Union, assert_never, cast, get_args, get_origin
 
 import cloudpickle
@@ -398,23 +397,18 @@ def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
                 self.__egg_decls_thunk__,
                 Thunk.value(TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.ident, name)))),
             )
+        bound = self.__egg_tp__.to_just() if self.__egg_tp__.args else None
         if name in cls_decl.class_methods:
             return RuntimeFunction(
-                self.__egg_decls_thunk__,
-                Thunk.value(ClassMethodRef(self.__egg_tp__.ident, name)),
-                self.__egg_tp__.to_just(),
+                self.__egg_decls_thunk__, Thunk.value(ClassMethodRef(self.__egg_tp__.ident, name)), bound
             )
         # allow referencing properties and methods as class variables as well
         if name in cls_decl.properties:
             return RuntimeFunction(
-                self.__egg_decls_thunk__,
-                Thunk.value(PropertyRef(self.__egg_tp__.ident, name)),
-                self.__egg_tp__.to_just(),
+                self.__egg_decls_thunk__, Thunk.value(PropertyRef(self.__egg_tp__.ident, name)), bound
             )
         if name in cls_decl.methods:
-            return RuntimeFunction(
-                self.__egg_decls_thunk__, Thunk.value(MethodRef(self.__egg_tp__.ident, name)), self.__egg_tp__.to_just()
-            )
+            return RuntimeFunction(self.__egg_decls_thunk__, Thunk.value(MethodRef(self.__egg_tp__.ident, name)), bound)
 
         msg = f"Class {self.__egg_tp__.ident} has no method {name}"
         raise AttributeError(msg) from None
@@ -504,7 +498,7 @@ def __hash__(self) -> int:
     def __egg_ref__(self) -> CallableRef:
         return self.__egg_ref_thunk__()
 
-    def __call__(  # noqa: C901
+    def __call__(  # noqa: C901,PLR0912
         self, *args: object, _egg_function_types: tuple[TypeOrVarRef, ...] | None = None, **kwargs: object
     ) -> RuntimeExpr | None:
         from .conversion import resolve_literal  # noqa: PLC0415
@@ -551,20 +545,14 @@ def __call__(  # noqa: C901
         args = bound.args
 
         tcs = TypeConstraintSolver()
-        bound_tp = (
-            None
-            if self.__egg_bound__ is None
-            else self.__egg_bound__.__egg_typed_expr__.tp
-            if isinstance(self.__egg_bound__, RuntimeExpr)
-            else self.__egg_bound__
-        )
-        if (
-            bound_tp
-            and bound_tp.args
-            # Don't  bind class if we have a first class function arg, b/c we don't support that yet
-            and not function_value
-        ):
-            tcs.bind_class(bound_tp, decls)
+        if isinstance(self.__egg_bound__, JustTypeRef) and self.__egg_bound__.args:
+            if function_value:
+                msg = "Cannot have both bound type params and function value"
+                raise ValueError(msg)
+            tcs.bind_class(self.__egg_bound__, decls)
+            bound_tp_params = self.__egg_bound__.args
+        else:
+            bound_tp_params = ()
         assert (operator.ge if signature.var_arg_type else operator.eq)(len(args), len(signature.arg_types))
         # Hack to allow being explicit on function types when casting. # noqa: FIX004
         for _fn_tp in _egg_function_types or ():
@@ -575,20 +563,29 @@ def __call__(  # noqa: C901
             tcs.bind_class(_fn_tp_just, decls)
             if _fn_tp_just.args:
                 pass
+        # Try using any runtime expressions passed in to help infer typevars
+        for arg, tp in zip(args, signature.all_args, strict=False):
+            if not isinstance(arg, RuntimeExpr):
+                continue
+            try:
+                tcs.infer_typevars(tp, arg.__egg_typed_expr__.tp)
+            # If this leads to an incompatibility, just skip it, since it could need to be upcasted
+            except TypeConstraintError:
+                continue
+        # Now at this point we should be able to resolve all the typevars
         upcasted_args = [
-            resolve_literal(cast("TypeOrVarRef", tp), arg, Thunk.value(decls), tcs)
-            for arg, tp in zip_longest(args, signature.arg_types, fillvalue=signature.var_arg_type)
+            resolve_literal(
+                tcs.substitute_typevars_try_function(tp, arg, Thunk.value(decls)).to_var(), arg, Thunk.value(decls)
+            )
+            for arg, tp in zip(args, signature.all_args, strict=False)
         ]
         decls.update(*upcasted_args)
         arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args)
         return_tp = tcs.substitute_typevars(signature.semantic_return_type)
-        bound_params = (
-            cast("JustTypeRef", bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef | InitRef) else ()
-        )
         # If we were using unstable-app to call a function, add that function back as the first arg.
         if function_value:
             arg_exprs = (function_value, *arg_exprs)
-        expr_decl = CallDecl(self.__egg_ref__, arg_exprs, bound_params)
+        expr_decl = CallDecl(self.__egg_ref__, arg_exprs, bound_tp_params)
         typed_expr_decl = TypedExprDecl(return_tp, expr_decl)
         # If there is not return type, we are mutating the first arg
         if not signature.return_type:
@@ -901,8 +898,7 @@ def create_callable(decls: Declarations, ref: CallableRef) -> RuntimeClass | Run
         case InitRef(name):
             return RuntimeClass(Thunk.value(decls), TypeRefWithVars(name))
         case FunctionRef() | MethodRef() | ClassMethodRef() | PropertyRef() | UnnamedFunctionRef():
-            bound = JustTypeRef(ref.ident) if isinstance(ref, ClassMethodRef) else None
-            return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), bound)
+            return RuntimeFunction(Thunk.value(decls), Thunk.value(ref), None)
         case ConstantRef(name):
             tp = decls._constants[name].type_ref
         case ClassVariableRef(cls_name, var_name):