egglog 7.0.0__cp312-none-win_amd64.whl → 7.1.0__cp312-none-win_amd64.whl

egglog/egraph_state.py

@@ -87,17 +87,26 @@ class EGraphState:
         """
         Registers a ruleset if it's not already registered.
         """
-        if name not in self.rulesets:
-            if name:
-                self.egraph.run_program(bindings.AddRuleset(name))
-            rules = self.rulesets[name] = set()
-        else:
-            rules = self.rulesets[name]
-        for rule in self.__egg_decls__._rulesets[name].rules:
-            if rule in rules:
-                continue
-            self.egraph.run_program(self.command_to_egg(rule, name))
-            rules.add(rule)
+        match self.__egg_decls__._rulesets[name]:
+            case RulesetDecl(rules):
+                if name not in self.rulesets:
+                    if name:
+                        self.egraph.run_program(bindings.AddRuleset(name))
+                    added_rules = self.rulesets[name] = set()
+                else:
+                    added_rules = self.rulesets[name]
+                for rule in rules:
+                    if rule in added_rules:
+                        continue
+                    self.egraph.run_program(self.command_to_egg(rule, name))
+                    added_rules.add(rule)
+            case CombinedRulesetDecl(rulesets):
+                if name in self.rulesets:
+                    return
+                self.rulesets[name] = set()
+                for ruleset in rulesets:
+                    self.ruleset_to_egg(ruleset)
+                self.egraph.run_program(bindings.UnstableCombinedRuleset(name, list(rulesets)))
 
     def command_to_egg(self, cmd: CommandDecl, ruleset: str) -> bindings._Command:
         match cmd:
@@ -184,11 +193,13 @@ class EGraphState:
                 )
             case FunctionDecl():
                 if not decl.builtin:
+                    signature = decl.signature
+                    assert isinstance(signature, FunctionSignature), "Cannot turn special function to egg"
                     egg_fn_decl = bindings.FunctionDecl(
                         egg_name,
                         bindings.Schema(
-                            [self.type_ref_to_egg(a.to_just()) for a in decl.arg_types],
-                            self.type_ref_to_egg(decl.semantic_return_type.to_just()),
+                            [self.type_ref_to_egg(a.to_just()) for a in signature.arg_types],
+                            self.type_ref_to_egg(signature.semantic_return_type.to_just()),
                         ),
                         self.expr_to_egg(decl.default) if decl.default else None,
                         self.expr_to_egg(decl.merge) if decl.merge else None,
@@ -212,19 +223,22 @@ class EGraphState:
         decl = self.__egg_decls__._classes[ref.name]
         self.type_ref_to_egg_sort[ref] = egg_name = decl.egg_name or _generate_type_egg_name(ref)
         if not decl.builtin or ref.args:
-            self.egraph.run_program(
-                bindings.Sort(
-                    egg_name,
-                    (
-                        (
-                            self.type_ref_to_egg(JustTypeRef(ref.name)),
-                            [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args],
-                        )
-                        if ref.args
-                        else None
-                    ),
-                )
-            )
+            if ref.args:
+                if ref.name == "UnstableFn":
+                    # UnstableFn is a special case, where the rest of args are collected into a call
+                    type_args: list[bindings._Expr] = [
+                        bindings.Call(
+                            self.type_ref_to_egg(ref.args[1]),
+                            [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args[2:]],
+                        ),
+                        bindings.Var(self.type_ref_to_egg(ref.args[0])),
+                    ]
+                else:
+                    type_args = [bindings.Var(self.type_ref_to_egg(a)) for a in ref.args]
+                args = (self.type_ref_to_egg(JustTypeRef(ref.name)), type_args)
+            else:
+                args = None
+            self.egraph.run_program(bindings.Sort(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.
@@ -292,6 +306,9 @@ class EGraphState:
                 res = bindings.Call(egg_fn, egg_args)
             case PyObjectDecl(value):
                 res = GLOBAL_PY_OBJECT_SORT.store(value)
+            case PartialCallDecl(call_decl):
+                egg_fn_call = self.expr_to_egg(call_decl)
+                res = bindings.Call("unstable-fn", [bindings.Lit(bindings.String(egg_fn_call.name)), *egg_fn_call.args])
             case _:
                 assert_never(expr_decl.expr)
 
@@ -371,24 +388,45 @@ class FromEggState:
             if term.name == "py-object":
                 call = bindings.termdag_term_to_expr(self.termdag, term)
                 expr_decl = PyObjectDecl(self.state.egraph.eval_py_object(call))
+            if term.name == "unstable-fn":
+                # Get function name
+                fn_term, *arg_terms = term.args
+                fn_value = self.resolve_term(fn_term, JustTypeRef("String"))
+                assert isinstance(fn_value.expr, LitDecl)
+                fn_name = fn_value.expr.value
+                assert isinstance(fn_name, str)
+
+                # Resolve what types the partiallied applied args are
+                assert tp.name == "UnstableFn"
+                call_decl = self.from_call(tp.args[0], bindings.TermApp(fn_name, arg_terms))
+                expr_decl = PartialCallDecl(call_decl)
             else:
                 expr_decl = self.from_call(tp, term)
         else:
             assert_never(term)
         return TypedExprDecl(tp, expr_decl)
 
-    def from_call(self, tp: JustTypeRef, term: bindings.TermApp) -> CallDecl:
+    def from_call(
+        self,
+        tp: JustTypeRef,
+        term: bindings.TermApp,  # additional_arg_tps: tuple[JustTypeRef, ...]
+    ) -> CallDecl:
         """
         Convert a call to a CallDecl.
 
         There could be Python call refs which match the call, so we need to find the correct one.
+
+        The additional_arg_tps are known types for arguments that come after the term args, used to infer types
+        for partially applied functions, where we know the types of the later args, but not of the earlier ones where
+        we have values for.
         """
         # Find the first callable ref that matches the call
         for callable_ref in self.state.egg_fn_to_callable_refs[term.name]:
             # If this is a classmethod, we might need the type params that were bound for this type
             # This could be multiple types if the classmethod is ambiguous, like map create.
             possible_types: Iterable[JustTypeRef | None]
-            fn_decl = self.decls.get_callable_decl(callable_ref).to_function_decl()
+            signature = self.decls.get_callable_decl(callable_ref).to_function_decl().signature
+            assert isinstance(signature, FunctionSignature)
             if isinstance(callable_ref, ClassMethodRef):
                 possible_types = self.state._get_possible_types(callable_ref.class_name)
                 cls_name = callable_ref.class_name
@@ -402,16 +440,17 @@ class FromEggState:
 
                 try:
                     arg_types, bound_tp_params = tcs.infer_arg_types(
-                        fn_decl.arg_types, fn_decl.semantic_return_type, fn_decl.var_arg_type, tp, cls_name
+                        signature.arg_types, signature.semantic_return_type, signature.var_arg_type, tp, cls_name
                     )
                 except TypeConstraintError:
                     continue
-                args: list[TypedExprDecl] = []
-                for a, tp in zip(term.args, arg_types, strict=False):
-                    try:
-                        res = self.cache[a]
-                    except KeyError:
-                        res = self.cache[a] = self.from_expr(tp, self.termdag.nodes[a])
-                    args.append(res)
-                return CallDecl(callable_ref, tuple(args), bound_tp_params)
+                args = tuple(self.resolve_term(a, tp) for a, tp in zip(term.args, arg_types, strict=False))
+                return CallDecl(callable_ref, args, bound_tp_params)
         raise ValueError(f"Could not find callable ref for call {term}")
+
+    def resolve_term(self, term_id: int, tp: JustTypeRef) -> TypedExprDecl:
+        try:
+            return self.cache[term_id]
+        except KeyError:
+            res = self.cache[term_id] = self.from_expr(tp, self.termdag.nodes[term_id])
+            return res

egglog/exp/array_api.py

@@ -18,7 +18,7 @@ from egglog.runtime import RuntimeExpr
 from .program_gen import *
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator
     from types import ModuleType
 
 # Pretend that exprs are numbers b/c sklearn does isinstance checks
@@ -257,7 +257,7 @@ class TupleInt(Expr):
 
     def __getitem__(self, i: Int) -> Int: ...
 
-    def product(self) -> Int: ...
+    def fold(self, init: Int, f: Callable[[Int, Int], Int]) -> Int: ...
 
 
 converter(
@@ -272,7 +272,7 @@ converter(
 
 
 @array_api_ruleset.register
-def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64):
+def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64, f: Callable[[Int, Int], Int]):
     return [
         rewrite(ti + TupleInt.EMPTY).to(ti),
         rewrite(TupleInt(i).length()).to(Int(1)),
@@ -281,10 +281,10 @@ def _tuple_int(ti: TupleInt, ti2: TupleInt, i: Int, i2: Int, k: i64):
         rewrite((TupleInt(i) + ti)[Int(0)]).to(i),
         # Rule for indexing > 0
         rule(eq(i).to((TupleInt(i2) + ti)[Int(k)]), k > 0).then(union(i).with_(ti[Int(k - 1)])),
-        # Product
-        rewrite(TupleInt(i).product()).to(i),
-        rewrite((TupleInt(i) + ti).product()).to(i * ti.product()),
-        rewrite(TupleInt.EMPTY.product()).to(Int(1)),
+        # fold
+        rewrite(TupleInt.EMPTY.fold(i, f)).to(i),
+        rewrite(TupleInt(i2).fold(i, f)).to(f(i, i2)),
+        rewrite((TupleInt(i2) + ti).fold(i, f)).to(ti.fold(f(i, i2), f)),
     ]
 
 
@@ -1346,7 +1346,7 @@ def _unique(xs: TupleValue, a: NDArray, shape: TupleInt, copy: OptionalBool):
 
 @array_api_ruleset.register
 def _size(x: NDArray):
-    yield rewrite(x.size).to(x.shape.product())
+    yield rewrite(x.size).to(x.shape.fold(Int(1), Int.__mul__))
 
 
 @overload

egglog/pretty.py

@@ -66,7 +66,7 @@ UNARY_METHODS = {
     "__invert__": "~",
 }
 
-AllDecls: TypeAlias = RulesetDecl | CommandDecl | ActionDecl | FactDecl | ExprDecl | ScheduleDecl
+AllDecls: TypeAlias = RulesetDecl | CombinedRulesetDecl | CommandDecl | ActionDecl | FactDecl | ExprDecl | ScheduleDecl
 
 
 def pretty_decl(
@@ -106,7 +106,7 @@ def pretty_callable_ref(
     """
     # Pass in three dummy args, which are the max used for any operation that
     # is not a generic function call
-    args: list[ExprDecl] = [LitDecl(ARG_STR)] * 3
+    args: list[ExprDecl] = [VarDecl(ARG_STR)] * 3
     if first_arg:
         args.insert(0, first_arg)
     res = PrettyContext(decls, defaultdict(lambda: 0))._call_inner(
@@ -117,6 +117,10 @@ def pretty_callable_ref(
     return res[0] if isinstance(res, tuple) else res
 
 
+# TODO: Add a different pretty callable ref that doesnt fill in wholes but instead returns the function
+# so that things like Math.__add__ will be represented properly
+
+
 @dataclass
 class TraverseContext:
     """
@@ -170,6 +174,10 @@ class TraverseContext:
                 if until:
                     for f in until:
                         self(f)
+            case PartialCallDecl(c):
+                self(c)
+            case CombinedRulesetDecl(_):
+                pass
             case _:
                 assert_never(decl)
 
@@ -231,6 +239,9 @@ class PrettyContext:
                 return name, name
             case CallDecl(_, _, _):
                 return self._call(decl, parens)
+            case PartialCallDecl(CallDecl(ref, typed_args, _)):
+                arg_strs = (_pretty_callable(ref), *(self(a.expr, parens=False, unwrap_lit=True) for a in typed_args))
+                return f"UnstableFn({', '.join(arg_strs)})", "fn"
             case PyObjectDecl(value):
                 return repr(value) if unwrap_lit else f"PyObject({value!r})", "PyObject"
             case ActionCommandDecl(action):
@@ -267,6 +278,10 @@ class PrettyContext:
                     return f"ruleset(name={ruleset_name!r})", f"ruleset_{ruleset_name}"
                 args = ", ".join(map(self, rules))
                 return f"ruleset({args})", "ruleset"
+            case CombinedRulesetDecl(rulesets):
+                if ruleset_name:
+                    rulesets = (*rulesets, f"name={ruleset_name!r})")
+                return f"unstable_combine_rulesets({', '.join(rulesets)})", "combined_ruleset"
             case SaturateDecl(schedule):
                 return f"{self(schedule, parens=True)}.saturate()", "schedule"
             case RepeatDecl(schedule, times):
@@ -302,19 +317,28 @@ class PrettyContext:
             l, r = self(args[0]), self(args[1])
             return f"ne({l}).to({r})", "Unit"
         function_decl = self.decls.get_callable_decl(ref).to_function_decl()
+        signature = function_decl.signature
+
         # Determine how many of the last arguments are defaults, by iterating from the end and comparing the arg with the default
         n_defaults = 0
-        for arg, default in zip(
-            reversed(args), reversed(function_decl.arg_defaults), strict=not function_decl.var_arg_type
-        ):
-            if arg != default:
-                break
-            n_defaults += 1
+        # Dont try counting defaults for function application
+        if isinstance(signature, FunctionSignature):
+            for arg, default in zip(
+                reversed(args), reversed(signature.arg_defaults), strict=not signature.var_arg_type
+            ):
+                if arg != default:
+                    break
+                n_defaults += 1
         if n_defaults:
             args = args[:-n_defaults]
 
-        tp_name = function_decl.semantic_return_type.name
-        if function_decl.mutates:
+        # If this is a function application, the type is the first type arg of the function object
+        if signature == "fn-app":
+            tp_name = decl.args[0].tp.args[0].name
+        else:
+            assert isinstance(signature, FunctionSignature)
+            tp_name = signature.semantic_return_type.name
+        if isinstance(signature, FunctionSignature) and signature.mutates:
             first_arg = args[0]
             expr_str = self(first_arg)
             # copy an identifier expression iff it has multiple parents (b/c then we can't mutate it directly)
@@ -397,6 +421,28 @@ class PrettyContext:
         return name
 
 
+def _pretty_callable(ref: CallableRef) -> str:
+    """
+    Returns a function call as a string.
+    """
+    match ref:
+        case FunctionRef(name):
+            return name
+        case (
+            ClassMethodRef(class_name, method_name)
+            | MethodRef(class_name, method_name)
+            | PropertyRef(class_name, method_name)
+        ):
+            return f"{class_name}.{method_name}"
+        case ConstantRef(_):
+            msg = "Constants should not be callable"
+            raise NotImplementedError(msg)
+        case ClassVariableRef(_, _):
+            msg = "Class variables should not be callable"
+            raise NotADirectoryError(msg)
+    assert_never(ref)
+
+
 def _plot_line_length(expr: object):  # pragma: no cover
     """
     Plots the number of line lengths based on different max lengths

egglog/runtime.py

@@ -11,7 +11,8 @@ so they are not mangled by Python and can be accessed by the user.
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+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
@@ -22,7 +23,7 @@ from .thunk import Thunk
 from .type_constraint_solver import *
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable
+    from collections.abc import Iterable
 
     from .egraph import Expr
 
@@ -60,6 +61,8 @@ REFLECTED_BINARY_METHODS = {
 # 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
+# Same for functions
+_UNSTABLE_FN_CLASS: RuntimeClass | None = None
 
 T = TypeVar("T")
 
@@ -67,6 +70,8 @@ T = TypeVar("T")
 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__)
@@ -79,6 +84,11 @@ 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__
@@ -95,9 +105,11 @@ class RuntimeClass(DelayedDeclerations):
     __egg_tp__: TypeRefWithVars
 
     def __post_init__(self) -> None:
-        global _PY_OBJECT_CLASS
-        if self.__egg_tp__.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
 
     def verify(self) -> None:
         if not self.__egg_tp__.args:
@@ -113,26 +125,48 @@ class RuntimeClass(DelayedDeclerations):
         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_tp__.name == "PyObject":
+        if (name := self.__egg_tp__.name) == "PyObject":
             assert len(args) == 1
             return RuntimeExpr.__from_value__(
                 self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), PyObjectDecl(args[0]))
             )
-        if self.__egg_tp__.name in UNARY_LIT_CLASS_NAMES:
+        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.__from_value__(
                 self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(args[0]))
             )
-        if self.__egg_tp__.name == UNIT_CLASS_NAME:
+        if name == UNIT_CLASS_NAME:
             assert len(args) == 0
             return RuntimeExpr.__from_value__(
                 self.__egg_decls__, TypedExprDecl(self.__egg_tp__.to_just(), LitDecl(None))
             )
 
         return RuntimeFunction(
-            Thunk.value(self.__egg_decls__), ClassMethodRef(self.__egg_tp__.name, "__init__"), self.__egg_tp__.to_just()
-        )(*args, **kwargs)
+            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_tp__.name)
@@ -184,6 +218,12 @@ class RuntimeClass(DelayedDeclerations):
             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(...)?"
@@ -207,24 +247,47 @@ class RuntimeFunction(DelayedDeclerations):
     # 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
 
-    def __call__(self, *args: object, **kwargs: object) -> RuntimeExpr | None:
+    def __call__(self, *args: object, _egg_partial_function: bool = False, **kwargs: object) -> RuntimeExpr | None:
         from .conversion import resolve_literal
 
         if isinstance(self.__egg_bound__, RuntimeExpr):
             args = (self.__egg_bound__, *args)
-        fn_decl = self.__egg_decls__.get_callable_decl(self.__egg_ref__).to_function_decl()
+        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
-        bound = callable_decl_to_signature(fn_decl, self.__egg_decls__).bind(*args, **kwargs)
+        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
-        del args, kwargs
+        args = bound.args
 
         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)
+            for arg, tp in zip_longest(args, signature.arg_types, fillvalue=signature.var_arg_type)
         ]
-
-        decls = Declarations.create(self, *upcasted_args)
+        decls.update(*upcasted_args)
 
         tcs = TypeConstraintSolver(decls)
         bound_tp = (
@@ -234,19 +297,27 @@ class RuntimeFunction(DelayedDeclerations):
             if isinstance(self.__egg_bound__, RuntimeExpr)
             else self.__egg_bound__
         )
-        if bound_tp and bound_tp.args:
+        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_tp.name if bound_tp else None
         return_tp = tcs.infer_return_type(
-            fn_decl.arg_types, fn_decl.return_type or fn_decl.arg_types[0], fn_decl.var_arg_type, arg_types, cls_name
+            signature.arg_types, signature.semantic_return_type, signature.var_arg_type, arg_types, cls_name
         )
         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 fn_decl.return_type:
+        if not signature.return_type:
             first_arg = upcasted_args[0]
             first_arg.__egg_thunk__ = Thunk.value((decls, typed_expr_decl))
             return None
@@ -262,19 +333,26 @@ class RuntimeFunction(DelayedDeclerations):
         return pretty_callable_ref(self.__egg_decls__, self.__egg_ref__, first_arg, bound_tp_params)
 
 
-def callable_decl_to_signature(
-    decl: FunctionDecl,
-    decls: Declarations,
-) -> Signature:
+def to_py_signature(sig: FunctionSignature, decls: Declarations, optional_args: bool) -> Signature:
+    """
+    Convert to a Python signature.
+
+    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.
+
+    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 else Parameter.empty,
+            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(decl.arg_names, decl.arg_defaults, decl.arg_types, strict=True)
+        for n, d, t in zip(sig.arg_names, sig.arg_defaults, sig.arg_types, strict=True)
     ]
-    if isinstance(decl, FunctionDecl) and decl.var_arg_type is not None:
+    if isinstance(sig, FunctionSignature) and sig.var_arg_type is not None:
         parameters.append(Parameter("__rest", Parameter.VAR_POSITIONAL))
     return Signature(parameters)
 
@@ -412,10 +490,14 @@ 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
+                # 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)
+            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)
@@ -436,8 +518,8 @@ def call_method_min_conversion(slf: object, other: object, name: str) -> Runtime
     # 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)
+    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)
 

egglog/thunk.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Generic, TypeVar
 
-from typing_extensions import ParamSpec, TypeVarTuple, Unpack
+from typing_extensions import TypeVarTuple, Unpack
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -12,7 +12,6 @@ if TYPE_CHECKING:
 __all__ = ["Thunk"]
 
 T = TypeVar("T")
-P = ParamSpec("P")
 TS = TypeVarTuple("TS")