Fix conversion to/from generic types

Author: saulshanabrook

docs/changelog.md

@@ -12,7 +12,7 @@ _This project uses semantic versioning_
 - Add working loopnest test and rewrite array api suport to be more general
 - Improve tracebacks on failing conversions.
 - Use `add_note` for exception to add more context, instead of raising a new exception, to make it easier to debug.
-- Add conversions from vecs (TODO)
+- Add conversions from generic types to be supported at runtime and typing level (so can go from `(1, 2, 3)` to `TupleInt`)
 - Open files with webbrowser instead of internal graphviz util for better support
 
 ## 8.0.1 (2024-10-24)

python/egglog/__init__.py

@@ -4,7 +4,7 @@
 
 from . import config, ipython_magic  # noqa: F401
 from .builtins import *  # noqa: UP029
-from .conversion import convert, converter  # noqa: F401
+from .conversion import ConvertError, convert, converter, get_type_args  # noqa: F401
 from .egraph import *
 
 del ipython_magic

python/egglog/builtins.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from functools import partial
+from functools import partial, reduce
 from types import FunctionType
 from typing import TYPE_CHECKING, Generic, Protocol, TypeAlias, TypeVar, Union, cast, overload
 
@@ -25,13 +25,16 @@
     "Bool",
     "BoolLike",
     "Map",
+    "MapLike",
     "PyObject",
     "Rational",
     "Set",
+    "SetLike",
     "String",
     "StringLike",
     "UnstableFn",
     "Vec",
+    "VecLike",
     "f64",
     "f64Like",
     "i64",
@@ -285,6 +288,22 @@ def remove(self, key: T) -> Map[T, V]: ...
     def rebuild(self) -> Map[T, V]: ...
 
 
+TO = TypeVar("TO")
+VO = TypeVar("VO")
+
+converter(
+    dict,
+    Map,
+    lambda t: reduce(
+        (lambda acc, kv: acc.insert(convert(kv[0], get_type_args()[0]), convert(kv[1], get_type_args()[1]))),
+        t.items(),
+        Map[get_type_args()].empty(),  # type: ignore[misc]
+    ),
+)
+
+MapLike: TypeAlias = Map[T, V] | dict[TO, VO]
+
+
 class Set(Expr, Generic[T], builtin=True):
     @method(egg_fn="set-of")
     def __init__(self, *args: T) -> None: ...
@@ -318,6 +337,17 @@ def __and__(self, other: Set[T]) -> Set[T]: ...
     def rebuild(self) -> Set[T]: ...
 
 
+converter(
+    set,
+    Set,
+    lambda t: Set[get_type_args()[0]](  # type: ignore[misc,operator]
+        *(convert(x, get_type_args()[0]) for x in t)
+    ),
+)
+
+SetLike: TypeAlias = Set[T] | set[TO]
+
+
 class Rational(Expr, builtin=True):
     @method(egg_fn="rational")
     def __init__(self, num: i64Like, den: i64Like) -> None: ...
@@ -418,7 +448,16 @@ def remove(self, index: i64Like) -> Vec[T]: ...
     def set(self, index: i64Like, value: T) -> Vec[T]: ...
 
 
-converter(tuple, Vec, lambda t: Vec(*(convert(x, get_type_args()[0]) for x in t)))
+for sequence_type in (list, tuple):
+    converter(
+        sequence_type,
+        Vec,
+        lambda t: Vec[get_type_args()[0]](  # type: ignore[misc,operator]
+            *(convert(x, get_type_args()[0]) for x in t)
+        ),
+    )
+
+VecLike: TypeAlias = Vec[T] | tuple[TO, ...] | list[TO]
 
 
 class PyObject(Expr, builtin=True):

python/egglog/conversion.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+from collections import defaultdict
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, NewType, TypeVar, cast
+from typing import TYPE_CHECKING, TypeVar, cast
 
 from .declarations import *
 from .pretty import *
@@ -15,10 +16,9 @@
 
     from .egraph import Expr
 
-__all__ = ["convert", "convert_to_same_type", "converter", "resolve_literal"]
+__all__ = ["convert", "convert_to_same_type", "converter", "resolve_literal", "ConvertError"]
 # 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]] = {}
+CONVERSIONS: dict[tuple[type | JustTypeRef, JustTypeRef], tuple[int, Callable]] = {}
 # 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
@@ -45,12 +45,12 @@ def converter(from_type: type[T], to_type: type[V], fn: Callable[[T], V], cost:
     Register a converter from some type to an egglog type.
     """
     to_type_name = process_tp(to_type)
-    if not isinstance(to_type_name, str):
+    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)
 
 
-def _register_converter(a: type | TypeName, b: TypeName, a_b: Callable, cost: int) -> None:
+def _register_converter(a: type | JustTypeRef, b: JustTypeRef, a_b: Callable, cost: int) -> None:
     """
     Registers a converter from some type to an egglog type, if not already registered.
 
@@ -63,10 +63,26 @@ def _register_converter(a: type | TypeName, b: TypeName, a_b: Callable, cost: in
         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)
+        if _is_type_compatible(b, c):
+            _register_converter(
+                a, d, _ComposedConverter(a_b, c_d, c.args if isinstance(c, JustTypeRef) else ()), cost + other_cost
+            )
+        if _is_type_compatible(a, d):
+            _register_converter(
+                c, b, _ComposedConverter(c_d, a_b, a.args if isinstance(a, JustTypeRef) else ()), cost + other_cost
+            )
+
+
+def _is_type_compatible(source: type | JustTypeRef, target: type | JustTypeRef) -> bool:
+    """
+    Types must be equal or also support unbound to bound typevar like B -> B[C]
+    """
+    if source == target:
+        return True
+    if isinstance(source, JustTypeRef) and isinstance(target, JustTypeRef) and source.args and not target.args:
+        return source.name == target.name
+        # TODO: Support case where B[T] where T is typevar is mapped to B[C]
+    return False
 
 
 @dataclass
@@ -81,9 +97,17 @@ class _ComposedConverter:
 
     a_b: Callable
     b_c: Callable
+    b_args: tuple[JustTypeRef, ...]
 
     def __call__(self, x: object) -> object:
-        return self.b_c(self.a_b(x))
+        # 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):
+                first_res = self.a_b(x)
+        else:
+            first_res = self.a_b(x)
+        return self.b_c(first_res)
 
     def __str__(self) -> str:
         return f"{self.b_c} ∘ {self.a_b}"
@@ -105,35 +129,33 @@ def convert_to_same_type(source: object, target: RuntimeExpr) -> RuntimeExpr:
     return resolve_literal(tp.to_var(), source, Thunk.value(target.__egg_decls__))
 
 
-def process_tp(tp: type | RuntimeClass) -> TypeName | type:
+def process_tp(tp: type | RuntimeClass) -> JustTypeRef | type:
     """
     Process a type before converting it, to add it to the global declerations and resolve to a ref.
     """
     if isinstance(tp, RuntimeClass):
         _TO_PROCESS_DECLS.append(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 egg_tp.to_just()
     return tp
 
 
-def min_convertable_tp(a: object, b: object, name: str) -> TypeName:
+def min_convertable_tp(a: object, b: object, name: str) -> JustTypeRef:
     """
     Returns the minimum convertable type between a and b, that has a method `name`, raising a ConvertError if no such type exists.
     """
     decls = _retrieve_conversion_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)
+        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)
+        to: c for ((from_, to), (c, _)) in CONVERSIONS.items() if from_ == b_tp and decls.has_method(to.name, name)
     }
-    if isinstance(a_tp, str):
+    if isinstance(a_tp, JustTypeRef):
         a_converts_to[a_tp] = 0
-    if isinstance(b_tp, str):
+    if isinstance(b_tp, JustTypeRef):
         b_converts_to[b_tp] = 0
     common = set(a_converts_to) & set(b_converts_to)
     if not common:
@@ -176,27 +198,38 @@ def resolve_literal(
         # If this is a var, it has to be a runtime expession
         assert isinstance(arg, RuntimeExpr), f"Expected a runtime expression, got {arg}"
         return arg
-    tp_name = TypeName(tp_just.name)
-    if arg_type == tp_name:
+    if arg_type == tp_just:
         # 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[(arg_type_instance, tp_name)][1]
-        except KeyError:
-            continue
-        break
+        if (key := (arg_type_instance, tp_just)) in CONVERSIONS:
+            fn = CONVERSIONS[key][1]
+            break
+        # Try broadening if we have a convert to the general type instead of the specific one too, for generics
+        if tp_just.args and (key := (arg_type_instance, JustTypeRef(tp_just.name))) in CONVERSIONS:
+            fn = CONVERSIONS[key][1]
+            break
+    # if we didn't find any raise an error
     else:
-        raise ConvertError(f"Cannot convert {arg_type} to {tp_name}")
+        raise ConvertError(f"Cannot convert {arg_type} to {tp_just}")
     with with_type_args(tp_just.args, decls):
         return fn(arg)
 
 
-def _get_tp(x: object) -> TypeName | type:
+def _debug_print_converers():
+    """
+    Prints a mapping of all source types to target types that have a conversion function.
+    """
+    source_to_targets = defaultdict(list)
+    for source, target in CONVERSIONS:
+        source_to_targets[source].append(target)
+
+
+def _get_tp(x: object) -> JustTypeRef | type:
     if isinstance(x, RuntimeExpr):
-        return TypeName(x.__egg_typed_expr__.tp.name)
+        return x.__egg_typed_expr__.tp
     tp = type(x)
     # If this value has a custom metaclass, let's use that as our index instead of the type
     if type(tp) is not type: