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

egglog/runtime.py

@@ -11,177 +11,67 @@ so they are not mangled by Python and can be accessed by the user.
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from inspect import Parameter, Signature
 from itertools import zip_longest
 from typing import TYPE_CHECKING, NoReturn, TypeVar, Union, cast, get_args, get_origin
 
-import black
-import black.parsing
-from typing_extensions import assert_never
-
-from . import bindings, config
 from .declarations import *
-from .declarations import BINARY_METHODS, REFLECTED_BINARY_METHODS, UNARY_METHODS
+from .pretty import *
+from .thunk import Thunk
 from .type_constraint_solver import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection, Iterable
+    from collections.abc import Iterable
 
     from .egraph import Expr
 
 __all__ = [
     "LIT_CLASS_NAMES",
-    "class_to_ref",
-    "resolve_literal",
     "resolve_callable",
     "resolve_type_annotation",
-    "convert_to_same_type",
     "RuntimeClass",
-    "RuntimeParamaterizedClass",
-    "RuntimeClassMethod",
     "RuntimeExpr",
     "RuntimeFunction",
-    "convert",
-    "converter",
+    "REFLECTED_BINARY_METHODS",
 ]
 
 
-BLACK_MODE = black.Mode(line_length=180)
-
 UNIT_CLASS_NAME = "Unit"
 UNARY_LIT_CLASS_NAMES = {"i64", "f64", "Bool", "String"}
 LIT_CLASS_NAMES = UNARY_LIT_CLASS_NAMES | {UNIT_CLASS_NAME, "PyObject"}
 
+REFLECTED_BINARY_METHODS = {
+    "__radd__": "__add__",
+    "__rsub__": "__sub__",
+    "__rmul__": "__mul__",
+    "__rmatmul__": "__matmul__",
+    "__rtruediv__": "__truediv__",
+    "__rfloordiv__": "__floordiv__",
+    "__rmod__": "__mod__",
+    "__rpow__": "__pow__",
+    "__rlshift__": "__lshift__",
+    "__rrshift__": "__rshift__",
+    "__rand__": "__and__",
+    "__rxor__": "__xor__",
+    "__ror__": "__or__",
+}
+
 # Set this globally so we can get access to PyObject when we have a type annotation of just object.
 # This is the only time a type annotation doesn't need to include the egglog type b/c object is top so that would be redundant statically.
 _PY_OBJECT_CLASS: RuntimeClass | None = None
-
-##
-# Converters
-##
-
-# 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]] = {}
-# Global declerations to store all convertable types so we can query if they have certain methods or not
-CONVERSIONS_DECLS = Declarations()
+# Same for functions
+_UNSTABLE_FN_CLASS: RuntimeClass | None = None
 
 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, 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 | JustTypeRef, b: JustTypeRef, 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.
-    """
-    target_ref = class_to_ref(cast(RuntimeTypeArgType, target))
-    return cast(V, resolve_literal(target_ref.to_var(), 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 | RuntimeTypeArgType) -> JustTypeRef | 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 | RuntimeParamaterizedClass):
-        CONVERSIONS_DECLS |= tp
-        return class_to_ref(tp)
-    return tp
-
-
-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 TypeError if no such type exists.
-    """
-    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 CONVERSIONS_DECLS.has_method(to.name, name)
-    }
-    b_converts_to = {
-        to: c
-        for ((from_, to), (c, _)) in CONVERSIONS.items()
-        if from_ == b_tp and CONVERSIONS_DECLS.has_method(to.name, name)
-    }
-    if isinstance(a_tp, JustTypeRef):
-        a_converts_to[a_tp] = 0
-    if isinstance(b_tp, JustTypeRef):
-        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_type_annotation(decls: Declarations, tp: object) -> TypeOrVarRef:
     """
     Resolves a type object into a type reference.
+
+    Any runtime type object decls will be add to those passed in.
     """
     if isinstance(tp, TypeVar):
         return ClassTypeVarRef(tp.__name__)
@@ -194,100 +84,92 @@ def resolve_type_annotation(decls: Declarations, tp: object) -> TypeOrVarRef:
     if tp == object:
         assert _PY_OBJECT_CLASS
         return resolve_type_annotation(decls, _PY_OBJECT_CLASS)
+    # If the type is a `Callable` then convert it into a UnstableFn
+    if get_origin(tp) == Callable:
+        assert _UNSTABLE_FN_CLASS
+        args, ret = get_args(tp)
+        return resolve_type_annotation(decls, _UNSTABLE_FN_CLASS[(ret, *args)])
     if isinstance(tp, RuntimeClass):
-        decls |= tp
-        return tp.__egg_tp__.to_var()
-    if isinstance(tp, RuntimeParamaterizedClass):
         decls |= tp
         return tp.__egg_tp__
     raise TypeError(f"Unexpected type annotation {tp}")
 
 
-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 exprssions
-        assert isinstance(arg, RuntimeExpr)
-        return arg
-    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[(cast(JustTypeRef | type, arg_type_instance), tp_just)][1]
-        except KeyError:
-            continue
-        break
-    else:
-        arg_type_str = arg_type.pretty() if isinstance(arg_type, JustTypeRef) else arg_type.__name__
-        raise ConvertError(f"Cannot convert {arg_type_str} to {tp_just.pretty()}")
-    return fn(arg)
-
-
-def _get_tp(x: object) -> JustTypeRef | type:
-    if isinstance(x, RuntimeExpr):
-        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) != type:
-        return type(tp)
-    return tp
-
-
 ##
 # Runtime objects
 ##
 
 
 @dataclass
-class RuntimeClass:
-    # Pass in a constructor to make the declarations lazy, so we can have classes reference each other in their type constructors
-    # This function should mutate the declerations and add to them
-    # Used this instead of a lazy property so we can have a reference to the decls in the class as its computing
-    lazy_decls: Callable[[Declarations], None] = field(repr=False)
-    # Cached declerations
-    _inner_decls: Declarations | None = field(init=False, repr=False, default=None)
-    __egg_name__: str
+class RuntimeClass(DelayedDeclerations):
+    __egg_tp__: TypeRefWithVars
 
     def __post_init__(self) -> None:
-        global _PY_OBJECT_CLASS
-        if self.__egg_name__ == "PyObject":
+        global _PY_OBJECT_CLASS, _UNSTABLE_FN_CLASS
+        if (name := self.__egg_tp__.name) == "PyObject":
             _PY_OBJECT_CLASS = self
+        elif name == "UnstableFn" and not self.__egg_tp__.args:
+            _UNSTABLE_FN_CLASS = self
 
-    @property
-    def __egg_decls__(self) -> Declarations:
-        if self._inner_decls is None:
-            # Set it like this so we can have a reference to the decls in the class as its computing
-            self._inner_decls = Declarations()
-            self.lazy_decls(self._inner_decls)
-        return self._inner_decls
+    def verify(self) -> None:
+        if not self.__egg_tp__.args:
+            return
+
+        # Raise error if we have args, but they are the wrong number
+        desired_args = self.__egg_decls__.get_class_decl(self.__egg_tp__.name).type_vars
+        if len(self.__egg_tp__.args) != len(desired_args):
+            raise ValueError(f"Expected {desired_args} type args, got {len(self.__egg_tp__.args)}")
 
     def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None:
         """
         Create an instance of this kind by calling the __init__ classmethod
         """
         # If this is a literal type, initializing it with a literal should return a literal
-        if self.__egg_name__ == "PyObject":
+        if (name := self.__egg_tp__.name) == "PyObject":
             assert len(args) == 1
-            return RuntimeExpr(self.__egg_decls__, TypedExprDecl(self.__egg_tp__, PyObjectDecl(args[0])))
-        if self.__egg_name__ in UNARY_LIT_CLASS_NAMES:
+            return RuntimeExpr.__from_value__(
+                self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), PyObjectDecl(args[0]))
+            )
+        if name == "UnstableFn":
+            assert not kwargs
+            fn_arg, *partial_args = args
+            del args
+            # Assumes we don't have types set for UnstableFn w/ generics, that they have to be inferred
+
+            # 1. Create a runtime function for the first arg
+            assert isinstance(fn_arg, RuntimeFunction)
+            # 2. Call it with the partial args, and use untyped vars for the rest of the args
+            res = fn_arg(*partial_args, _egg_partial_function=True)
+            assert res is not None, "Mutable partial functions not supported"
+            # 3. Use the inferred return type and inferred rest arg types as the types of the function, and
+            #    the partially applied args as the args.
+            call = (res_typed_expr := res.__egg_typed_expr__).expr
+            return_tp = res_typed_expr.tp
+            assert isinstance(call, CallDecl), "partial function must be a call"
+            n_args = len(partial_args)
+            value = PartialCallDecl(replace(call, args=call.args[:n_args]))
+            remaining_arg_types = [a.tp for a in call.args[n_args:]]
+            type_ref = JustTypeRef("UnstableFn", (return_tp, *remaining_arg_types))
+            return RuntimeExpr.__from_value__(Declarations.create(self, res), TypedExprDecl(type_ref, value))
+
+        if name in UNARY_LIT_CLASS_NAMES:
             assert len(args) == 1
             assert isinstance(args[0], int | float | str | bool)
-            return RuntimeExpr(self.__egg_decls__, TypedExprDecl(self.__egg_tp__, LitDecl(args[0])))
-        if self.__egg_name__ == UNIT_CLASS_NAME:
+            return RuntimeExpr.__from_value__(
+                self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(args[0]))
+            )
+        if name == UNIT_CLASS_NAME:
             assert len(args) == 0
-            return RuntimeExpr(self.__egg_decls__, TypedExprDecl(self.__egg_tp__, LitDecl(None)))
+            return RuntimeExpr.__from_value__(
+                self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(None))
+            )
 
-        return RuntimeClassMethod(self.__egg_decls__, self.__egg_tp__, "__init__")(*args, **kwargs)
+        return RuntimeFunction(
+            Thunk.value(self.__egg_decls__), ClassMethodRef(name, "__init__"), self.__egg_tp__.to_just()
+        )(*args, **kwargs)  # type: ignore[arg-type]
 
     def __dir__(self) -> list[str]:
-        cls_decl = self.__egg_decls__.get_class_decl(self.__egg_name__)
+        cls_decl = self.__egg_decls__.get_class_decl(self.__egg_tp__.name)
         possible_methods = (
             list(cls_decl.class_methods) + list(cls_decl.class_variables) + list(cls_decl.preserved_methods)
         )
@@ -296,14 +178,19 @@ class RuntimeClass:
             possible_methods.append("__call__")
         return possible_methods
 
-    def __getitem__(self, args: object) -> RuntimeParamaterizedClass:
+    def __getitem__(self, args: object) -> RuntimeClass:
+        if self.__egg_tp__.args:
+            raise TypeError(f"Cannot index into a paramaterized class {self}")
         if not isinstance(args, tuple):
             args = (args,)
         decls = self.__egg_decls__.copy()
-        tp = TypeRefWithVars(self.__egg_name__, tuple(resolve_type_annotation(decls, arg) for arg in args))
-        return RuntimeParamaterizedClass(self.__egg_decls__, tp)
+        tp = TypeRefWithVars(self.__egg_tp__.name, tuple(resolve_type_annotation(decls, arg) for arg in args))
+        return RuntimeClass(Thunk.value(decls), tp)
+
+    def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable:
+        if name == "__origin__" and self.__egg_tp__.args:
+            return RuntimeClass(self.__egg_decls_thunk__, TypeRefWithVars(self.__egg_tp__.name))
 
-    def __getattr__(self, name: str) -> RuntimeClassMethod | RuntimeExpr | Callable:
         # Special case some names that don't exist so we can exit early without resolving decls
         # Important so if we take union of RuntimeClass it won't try to resolve decls
         if name in {
@@ -314,7 +201,7 @@ class RuntimeClass:
         }:
             raise AttributeError
 
-        cls_decl = self.__egg_decls__.get_class_decl(self.__egg_name__)
+        cls_decl = self.__egg_decls__._classes[self.__egg_tp__.name]
 
         preserved_methods = cls_decl.preserved_methods
         if name in preserved_methods:
@@ -323,159 +210,151 @@ class RuntimeClass:
         # if this is a class variable, return an expr for it, otherwise, assume it's a method
         if name in cls_decl.class_variables:
             return_tp = cls_decl.class_variables[name]
-            return RuntimeExpr(
-                self.__egg_decls__, TypedExprDecl(return_tp, CallDecl(ClassVariableRef(self.__egg_name__, name)))
+            return RuntimeExpr.__from_value__(
+                self.__egg_decls__,
+                TypedExprDecl(return_tp.type_ref, CallDecl(ClassVariableRef(self.__egg_tp__.name, name))),
             )
-        return RuntimeClassMethod(self.__egg_decls__, self.__egg_tp__, name)
+        if name in cls_decl.class_methods:
+            return RuntimeFunction(
+                Thunk.value(self.__egg_decls__), ClassMethodRef(self.__egg_tp__.name, name), self.__egg_tp__.to_just()
+            )
+        # allow referencing properties and methods as class variables as well
+        if name in cls_decl.properties:
+            return RuntimeFunction(Thunk.value(self.__egg_decls__), PropertyRef(self.__egg_tp__.name, name))
+        if name in cls_decl.methods:
+            return RuntimeFunction(Thunk.value(self.__egg_decls__), MethodRef(self.__egg_tp__.name, name))
+
+        msg = f"Class {self.__egg_tp__.name} has no method {name}"
+        if name == "__ne__":
+            msg += ". Did you mean to use the ne(...).to(...)?"
+        raise AttributeError(msg) from None
 
     def __str__(self) -> str:
-        return self.__egg_name__
+        return str(self.__egg_tp__)
 
     # Make hashable so can go in Union
     def __hash__(self) -> int:
-        return hash((id(self.lazy_decls), self.__egg_name__))
+        return hash((id(self.__egg_decls_thunk__), self.__egg_tp__))
 
     # Support unioning like types
     def __or__(self, __value: type) -> object:
         return Union[self, __value]  # noqa: UP007
 
-    @property
-    def __egg_tp__(self) -> JustTypeRef:
-        return JustTypeRef(self.__egg_name__)
-
-
-@dataclass
-class RuntimeParamaterizedClass:
-    __egg_decls__: Declarations
-    __egg_tp__: TypeRefWithVars
-
-    def __post_init__(self) -> None:
-        desired_args = self.__egg_decls__.get_class_decl(self.__egg_tp__.name).type_vars
-        if len(self.__egg_tp__.args) != len(desired_args):
-            raise ValueError(f"Expected {desired_args} type args, got {len(self.__egg_tp__.args)}")
-
-    def __call__(self, *args: object) -> RuntimeExpr | None:
-        return RuntimeClassMethod(self.__egg_decls__, class_to_ref(self), "__init__")(*args)
-
-    def __getattr__(self, name: str) -> RuntimeClassMethod | RuntimeClass:
-        # Special case so when get_type_annotations proccessed it can work
-        if name == "__origin__":
-            return RuntimeClass(self.__egg_decls__.update_other, self.__egg_tp__.name)
-        return RuntimeClassMethod(self.__egg_decls__, class_to_ref(self), name)
-
-    def __str__(self) -> str:
-        return self.__egg_tp__.pretty()
-
-    # Support unioning
-    def __or__(self, __value: type) -> object:
-        return Union[self, __value]  # noqa: UP007
-
-
-# Type args can either be typevars or classes
-RuntimeTypeArgType = RuntimeClass | RuntimeParamaterizedClass
-
-
-def class_to_ref(cls: RuntimeTypeArgType) -> JustTypeRef:
-    if isinstance(cls, RuntimeClass):
-        return JustTypeRef(cls.__egg_name__)
-    if isinstance(cls, RuntimeParamaterizedClass):
-        # Currently this is used when calling methods on a parametrized class, which is only possible when we
-        # have actualy types currently, not typevars, currently.
-        return cls.__egg_tp__.to_just()
-    assert_never(cls)
-
 
 @dataclass
-class RuntimeFunction:
-    __egg_decls__: Declarations
-    __egg_name__: str
-    __egg_fn_ref__: FunctionRef = field(init=False)
-    __egg_fn_decl__: FunctionDecl = field(init=False)
-
-    def __post_init__(self) -> None:
-        self.__egg_fn_ref__ = FunctionRef(self.__egg_name__)
-        self.__egg_fn_decl__ = self.__egg_decls__.get_function_decl(self.__egg_fn_ref__)
-
-    def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None:
-        return _call(self.__egg_decls__, self.__egg_fn_ref__, self.__egg_fn_decl__, args, kwargs)
-
-    def __str__(self) -> str:
-        return self.__egg_name__
-
-
-def _call(
-    decls_from_fn: Declarations,
-    callable_ref: CallableRef,
-    fn_decl: FunctionDecl,
-    args: Collection[object],
-    kwargs: dict[str, object],
-    bound_class: JustTypeRef | None = None,
-) -> RuntimeExpr | None:
-    # Turn all keyword args into positional args
-    bound = fn_decl.to_signature(lambda expr: RuntimeExpr(decls_from_fn, expr)).bind(*args, **kwargs)
-    bound.apply_defaults()
-    assert not bound.kwargs
-    del args, kwargs
-
-    upcasted_args = [
-        resolve_literal(cast(TypeOrVarRef, tp), arg)
-        for arg, tp in zip_longest(bound.args, fn_decl.arg_types, fillvalue=fn_decl.var_arg_type)
-    ]
+class RuntimeFunction(DelayedDeclerations):
+    __egg_ref__: CallableRef
+    # bound methods need to store RuntimeExpr not just TypedExprDecl, so they can mutate the expr if required on self
+    __egg_bound__: JustTypeRef | RuntimeExpr | None = None
 
-    arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args)
-    decls = Declarations.create(decls_from_fn, *upcasted_args)
+    def __call__(self, *args: object, _egg_partial_function: bool = False, **kwargs: object) -> RuntimeExpr | None:
+        from .conversion import resolve_literal
 
-    tcs = TypeConstraintSolver(decls)
-    if bound_class is not None and bound_class.args:
-        tcs.bind_class(bound_class)
-
-    if fn_decl is not None:
+        if isinstance(self.__egg_bound__, RuntimeExpr):
+            args = (self.__egg_bound__, *args)
+        signature = self.__egg_decls__.get_callable_decl(self.__egg_ref__).to_function_decl().signature
+        decls = self.__egg_decls__.copy()
+        # Special case function application bc we dont support variadic generics yet generally
+        if signature == "fn-app":
+            fn, *rest_args = args
+            args = tuple(rest_args)
+            assert not kwargs
+            assert isinstance(fn, RuntimeExpr)
+            decls.update(fn)
+            function_value = fn.__egg_typed_expr__
+            fn_tp = function_value.tp
+            assert fn_tp.name == "UnstableFn"
+            fn_return_tp, *fn_arg_tps = fn_tp.args
+            signature = FunctionSignature(
+                tuple(tp.to_var() for tp in fn_arg_tps),
+                tuple(f"_{i}" for i in range(len(fn_arg_tps))),
+                (None,) * len(fn_arg_tps),
+                fn_return_tp.to_var(),
+            )
+        else:
+            function_value = None
+            assert isinstance(signature, FunctionSignature)
+
+        # Turn all keyword args into positional args
+        py_signature = to_py_signature(signature, self.__egg_decls__, _egg_partial_function)
+        bound = py_signature.bind(*args, **kwargs)
+        del kwargs
+        bound.apply_defaults()
+        assert not bound.kwargs
+        args = bound.args
+
+        upcasted_args = [
+            resolve_literal(cast(TypeOrVarRef, tp), arg)
+            for arg, tp in zip_longest(args, signature.arg_types, fillvalue=signature.var_arg_type)
+        ]
+        decls.update(*upcasted_args)
+
+        tcs = TypeConstraintSolver(decls)
+        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)
+        arg_exprs = tuple(arg.__egg_typed_expr__ for arg in upcasted_args)
         arg_types = [expr.tp for expr in arg_exprs]
-        cls_name = bound_class.name if bound_class is not None else None
+        cls_name = bound_tp.name if bound_tp else None
         return_tp = tcs.infer_return_type(
-            fn_decl.arg_types, fn_decl.return_type, fn_decl.var_arg_type, arg_types, cls_name
+            signature.arg_types, signature.semantic_return_type, signature.var_arg_type, arg_types, cls_name
         )
-    else:
-        return_tp = JustTypeRef("Unit")
-    bound_params = cast(JustTypeRef, bound_class).args if isinstance(callable_ref, ClassMethodRef) else None
-    expr_decl = CallDecl(callable_ref, arg_exprs, bound_params)
-    typed_expr_decl = TypedExprDecl(return_tp, expr_decl)
-    # Register return type sort in case it's a variadic generic that needs to be created
-    decls.register_sort(return_tp, False)
-    if fn_decl.mutates_first_arg:
-        first_arg = upcasted_args[0]
-        first_arg.__egg_typed_expr__ = typed_expr_decl
-        first_arg.__egg_decls__ = decls
-        return None
-    return RuntimeExpr(decls, TypedExprDecl(return_tp, expr_decl))
-
+        bound_params = cast(JustTypeRef, bound_tp).args if isinstance(self.__egg_ref__, ClassMethodRef) else None
+        # If we were using unstable-app to call a funciton, 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)
+        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:
+            first_arg = upcasted_args[0]
+            first_arg.__egg_thunk__ = Thunk.value((decls, typed_expr_decl))
+            return None
+        return RuntimeExpr.__from_value__(decls, typed_expr_decl)
 
-@dataclass
-class RuntimeClassMethod:
-    __egg_decls__: Declarations
-    __egg_tp__: JustTypeRef
-    __egg_method_name__: str
-    __egg_callable_ref__: ClassMethodRef = field(init=False)
-    __egg_fn_decl__: FunctionDecl = field(init=False)
+    def __str__(self) -> str:
+        first_arg, bound_tp_params = None, None
+        match self.__egg_bound__:
+            case RuntimeExpr(_):
+                first_arg = self.__egg_bound__.__egg_typed_expr__.expr
+            case JustTypeRef(_, args):
+                bound_tp_params = args
+        return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params)
 
-    def __post_init__(self) -> None:
-        self.__egg_callable_ref__ = ClassMethodRef(self.class_name, self.__egg_method_name__)
-        try:
-            self.__egg_fn_decl__ = self.__egg_decls__.get_function_decl(self.__egg_callable_ref__)
-        except KeyError as e:
-            raise AttributeError(f"Class {self.class_name} does not have method {self.__egg_method_name__}") from e
 
-    def __call__(self, *args: object, **kwargs) -> RuntimeExpr | None:
-        return _call(self.__egg_decls__, self.__egg_callable_ref__, self.__egg_fn_decl__, args, kwargs, self.__egg_tp__)
+def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: bool) -> Signature:
+    """
+    Convert to a Python signature.
 
-    def __str__(self) -> str:
-        return f"{self.class_name}.{self.__egg_method_name__}"
+    If optional_args is true, then all args will be treated as optional, as if a default was provided that makes them
+    a var with that arg name as the value.
 
-    @property
-    def class_name(self) -> str:
-        if isinstance(self.__egg_tp__, str):
-            return self.__egg_tp__
-        return self.__egg_tp__.name
+    Used for partial application to try binding a function with only some of its args.
+    """
+    parameters = [
+        Parameter(
+            n,
+            Parameter.POSITIONAL_OR_KEYWORD,
+            default=RuntimeExpr.__from_value__(decls, TypedExprDecl(t.to_just(), d if d is not None else VarDecl(n)))
+            if d is not None or optional_args
+            else Parameter.empty,
+        )
+        for n, d, t in zip(sig.arg_names, sig.arg_defaults, sig.arg_types, strict=True)
+    ]
+    if isinstance(sig, FunctionSignature) and sig.var_arg_type is not None:
+        parameters.append(Parameter("__rest", Parameter.VAR_POSITIONAL))
+    return Signature(parameters)
 
 
 # All methods which should return NotImplemented if they fail to resolve
@@ -505,63 +384,34 @@ PARTIAL_METHODS = {
 
 
 @dataclass
-class RuntimeMethod:
-    __egg_self__: RuntimeExpr
-    __egg_method_name__: str
-    __egg_callable_ref__: MethodRef | PropertyRef = field(init=False)
-    __egg_fn_decl__: FunctionDecl = field(init=False, repr=False)
-    __egg_decls__: Declarations = field(init=False)
-
-    def __post_init__(self) -> None:
-        self.__egg_decls__ = self.__egg_self__.__egg_decls__
-        if self.__egg_method_name__ in self.__egg_decls__.get_class_decl(self.class_name).properties:
-            self.__egg_callable_ref__ = PropertyRef(self.class_name, self.__egg_method_name__)
-        else:
-            self.__egg_callable_ref__ = MethodRef(self.class_name, self.__egg_method_name__)
-        try:
-            self.__egg_fn_decl__ = self.__egg_decls__.get_function_decl(self.__egg_callable_ref__)
-        except KeyError:
-            msg = f"Class {self.class_name} does not have method {self.__egg_method_name__}"
-            if self.__egg_method_name__ == "__ne__":
-                msg += ". Did you mean to use the ne(...).to(...)?"
-            raise AttributeError(msg) from None
+class RuntimeExpr:
+    # Defer needing decls/expr so we can make constants that don't resolve their class types
+    __egg_thunk__: Callable[[], tuple[Declarations, TypedExprDecl]]
 
-    def __call__(self, *args: object, **kwargs) -> RuntimeExpr | None:
-        args = (self.__egg_self__, *args)
-        try:
-            return _call(
-                self.__egg_decls__,
-                self.__egg_callable_ref__,
-                self.__egg_fn_decl__,
-                args,
-                kwargs,
-                self.__egg_self__.__egg_typed_expr__.tp,
-            )
-        except ConvertError as e:
-            name = self.__egg_method_name__
-            raise TypeError(f"Wrong types for {self.__egg_self__.__egg_typed_expr__.tp.pretty()}.{name}") from e
+    @classmethod
+    def __from_value__(cls, d: Declarations, e: TypedExprDecl) -> RuntimeExpr:
+        return cls(Thunk.value((d, e)))
 
     @property
-    def class_name(self) -> str:
-        return self.__egg_self__.__egg_typed_expr__.tp.name
-
+    def __egg_decls__(self) -> Declarations:
+        return self.__egg_thunk__()[0]
 
-@dataclass
-class RuntimeExpr:
-    __egg_decls__: Declarations
-    __egg_typed_expr__: TypedExprDecl
+    @property
+    def __egg_typed_expr__(self) -> TypedExprDecl:
+        return self.__egg_thunk__()[1]
 
-    def __getattr__(self, name: str) -> RuntimeMethod | RuntimeExpr | Callable | None:
-        class_decl = self.__egg_decls__.get_class_decl(self.__egg_typed_expr__.tp.name)
+    def __getattr__(self, name: str) -> RuntimeFunction | RuntimeExpr | Callable | None:
+        cls_name = self.__egg_class_name__
+        class_decl = self.__egg_class_decl__
 
-        preserved_methods = class_decl.preserved_methods
-        if name in preserved_methods:
+        if name in (preserved_methods := class_decl.preserved_methods):
             return preserved_methods[name].__get__(self)
 
-        method = RuntimeMethod(self, name)
-        if isinstance(method.__egg_callable_ref__, PropertyRef):
-            return method()
-        return method
+        if name in class_decl.methods:
+            return RuntimeFunction(Thunk.value(self.__egg_decls__), MethodRef(cls_name, name), self)
+        if name in class_decl.properties:
+            return RuntimeFunction(Thunk.value(self.__egg_decls__), PropertyRef(cls_name, name), self)()
+        raise AttributeError(f"{cls_name} has no method {name}") from None
 
     def __repr__(self) -> str:
         """
@@ -570,18 +420,10 @@ class RuntimeExpr:
         return str(self)
 
     def __str__(self) -> str:
-        context = PrettyContext(self.__egg_decls__)
-        context.traverse_for_parents(self.__egg_typed_expr__.expr)
-        pretty_expr = self.__egg_typed_expr__.expr.pretty(context, parens=False)
-        try:
-            if config.SHOW_TYPES:
-                raise NotImplementedError
-                # s = f"_: {self.__egg_typed_expr__.tp.pretty()} = {pretty_expr}"
-                # return black.format_str(s, mode=black.FileMode()).strip()
-            pretty_statements = context.render(pretty_expr)
-            return black.format_str(pretty_statements, mode=BLACK_MODE).strip()
-        except black.parsing.InvalidInput:
-            return pretty_expr
+        return self.__egg_pretty__(None)
+
+    def __egg_pretty__(self, wrapping_fn: str | None) -> str:
+        return pretty_decl(self.__egg_decls__, self.__egg_typed_expr__.expr, wrapping_fn=wrapping_fn)
 
     def _ipython_display_(self) -> None:
         from IPython.display import Code, display
@@ -589,28 +431,32 @@ class RuntimeExpr:
         display(Code(str(self), language="python"))
 
     def __dir__(self) -> Iterable[str]:
-        return list(self.__egg_decls__.get_class_decl(self.__egg_typed_expr__.tp.name).methods)
+        class_decl = self.__egg_class_decl__
+        return list(class_decl.methods) + list(class_decl.properties) + list(class_decl.preserved_methods)
 
     @property
-    def __egg__(self) -> bindings._Expr:
-        return self.__egg_typed_expr__.to_egg(self.__egg_decls__)
+    def __egg_class_name__(self) -> str:
+        return self.__egg_typed_expr__.tp.name
+
+    @property
+    def __egg_class_decl__(self) -> ClassDecl:
+        return self.__egg_decls__.get_class_decl(self.__egg_class_name__)
 
     # Have __eq__ take no NoReturn (aka Never https://docs.python.org/3/library/typing.html#typing.Never) because
     # we don't wany any type that MyPy thinks is an expr to be used with __eq__.
     # That's because we want to reserve __eq__ for domain specific equality checks, overloading this method.
     # To check if two exprs are equal, use the expr_eq method.
-    def __eq__(self, other: NoReturn) -> Expr:  # type: ignore[override]
-        msg = "Do not use == on RuntimeExpr. Compare the __egg_typed_expr__ attribute instead for structural equality."
-        raise NotImplementedError(msg)
+    # At runtime, this will resolve if there is a defined egg function for `__eq__`
+    def __eq__(self, other: NoReturn) -> Expr: ...  # type: ignore[override, empty-body]
 
     # Implement these so that copy() works on this object
     # otherwise copy will try to call `__getstate__` before object is initialized with properties which will cause inifinite recursion
 
     def __getstate__(self) -> tuple[Declarations, TypedExprDecl]:
-        return (self.__egg_decls__, self.__egg_typed_expr__)
+        return self.__egg_thunk__()
 
     def __setstate__(self, d: tuple[Declarations, TypedExprDecl]) -> None:
-        self.__egg_decls__, self.__egg_typed_expr__ = d
+        self.__egg_thunk__ = Thunk.value(d)
 
     def __hash__(self) -> int:
         return hash(self.__egg_typed_expr__)
@@ -625,12 +471,17 @@ for name in list(BINARY_METHODS) + list(UNARY_METHODS) + ["__getitem__", "__call
         __name: str = name,
         **kwargs: object,
     ) -> RuntimeExpr | None:
+        from .conversion import ConvertError
+
+        class_name = self.__egg_class_name__
+        class_decl = self.__egg_class_decl__
         # First, try to resolve as preserved method
         try:
-            method = self.__egg_decls__.get_class_decl(self.__egg_typed_expr__.tp.name).preserved_methods[__name]
-            return method(self, *args, **kwargs)
+            method = class_decl.preserved_methods[__name]
         except KeyError:
             pass
+        else:
+            return method(self, *args, **kwargs)
         # If this is a "partial" method meaning that it can return NotImplemented,
         # we want to find the "best" superparent (lowest cost) of the arg types to call with it, instead of just
         # using the arg type of the self arg.
@@ -639,8 +490,15 @@ for name in list(BINARY_METHODS) + list(UNARY_METHODS) + ["__getitem__", "__call
             try:
                 return call_method_min_conversion(self, args[0], __name)
             except ConvertError:
-                return NotImplemented
-        return RuntimeMethod(self, __name)(*args, **kwargs)
+                # Defer raising not imeplemented in case the dunder method is not symmetrical, then
+                # we use the standard process
+                pass
+        if __name in class_decl.methods:
+            fn = RuntimeFunction(Thunk.value(self.__egg_decls__), MethodRef(class_name, __name), self)
+            return fn(*args, **kwargs)  # type: ignore[arg-type]
+        if __name in PARTIAL_METHODS:
+            return NotImplemented
+        raise TypeError(f"{class_name!r} object does not support {__name}")
 
     setattr(RuntimeExpr, name, _special_method)
 
@@ -655,12 +513,14 @@ for reflected, non_reflected in REFLECTED_BINARY_METHODS.items():
 
 
 def call_method_min_conversion(slf: object, other: object, name: str) -> RuntimeExpr | None:
+    from .conversion import min_convertable_tp, resolve_literal
+
     # find a minimum type that both can be converted to
     # This is so so that calls like `-0.1 * Int("x")` work by upcasting both to floats.
     min_tp = min_convertable_tp(slf, other, name)
-    slf = resolve_literal(min_tp.to_var(), slf)
-    other = resolve_literal(min_tp.to_var(), other)
-    method = RuntimeMethod(slf, name)
+    slf = resolve_literal(TypeRefWithVars(min_tp), slf)
+    other = resolve_literal(TypeRefWithVars(min_tp), other)
+    method = RuntimeFunction(Thunk.value(slf.__egg_decls__), MethodRef(slf.__egg_class_name__, name), slf)
     return method(other)
 
 
@@ -680,21 +540,9 @@ def resolve_callable(callable: object) -> tuple[CallableRef, Declarations]:
     """
     Resolves a runtime callable into a ref
     """
-    # TODO: Fix these typings.
-    ref: CallableRef
-    decls: Declarations
-    if isinstance(callable, RuntimeFunction):
-        ref = FunctionRef(callable.__egg_name__)
-        decls = callable.__egg_decls__
-    elif isinstance(callable, RuntimeClassMethod):
-        ref = ClassMethodRef(callable.class_name, callable.__egg_method_name__)
-        decls = callable.__egg_decls__
-    elif isinstance(callable, RuntimeMethod):
-        ref = MethodRef(callable.__egg_self__.__egg_typed_expr__.tp.name, callable.__egg_method_name__)
-        decls = callable.__egg_decls__
-    elif isinstance(callable, RuntimeClass):
-        ref = ClassMethodRef(callable.__egg_name__, "__init__")
-        decls = callable.__egg_decls__
-    else:
-        raise NotImplementedError(f"Cannot turn {callable} into a callable ref")
-    return (ref, decls)
+    match callable:
+        case RuntimeFunction(decls, ref, _):
+            return ref, decls()
+        case RuntimeClass(thunk, tp):
+            return ClassMethodRef(tp.name, "__init__"), thunk()
+    raise NotImplementedError(f"Cannot turn {callable} into a callable ref")