Move unnamed functions to separate structure

Author: saulshanabrook

python/egglog/declarations.py

@@ -120,8 +120,8 @@ def upcast_declerations(declerations_like: Iterable[DeclerationsLike]) -> list[D
 
 @dataclass
 class Declarations:
-    # TODO: Replace with set of unnamed function decls
-    _functions: dict[str | UnnamedFunctionRef, FunctionDecl | RelationDecl] = field(default_factory=dict)
+    _unnamed_functions: set[UnnamedFunctionRef] = field(default_factory=set)
+    _functions: dict[str, FunctionDecl | RelationDecl] = field(default_factory=dict)
     _constants: dict[str, ConstantDecl] = field(default_factory=dict)
     _classes: dict[str, ClassDecl] = field(default_factory=dict)
     _rulesets: dict[str, RulesetDecl | CombinedRulesetDecl] = field(default_factory=lambda: {"": RulesetDecl([])})
@@ -197,6 +197,8 @@ def get_callable_decl(self, ref: CallableRef) -> CallableDecl:  # noqa: PLR0911
                 init_fn = self._classes[class_name].init
                 assert init_fn
                 return init_fn
+            case UnnamedFunctionRef():
+                return ref.to_function_decl()
         assert_never(ref)
 
     def set_function_decl(
@@ -329,20 +331,34 @@ class UnnamedFunctionRef:
     A reference to a function that doesn't have a name, but does have a body.
     """
 
-    arg_types: tuple[JustTypeRef, ...]
-    arg_names: tuple[str, ...]
+    # tuple of var arg names and their types
+    args: tuple[TypedExprDecl, ...]
     res: TypedExprDecl
 
-    @property
-    def args(self) -> tuple[TypedExprDecl, ...]:
-        return tuple(
-            TypedExprDecl(tp, VarDecl(name, False)) for tp, name in zip(self.arg_types, self.arg_names, strict=True)
+    def to_function_decl(self) -> FunctionDecl:
+        arg_types = []
+        arg_names = []
+        for a in self.args:
+            arg_types.append(a.tp.to_var())
+            assert isinstance(a.expr, VarDecl)
+            arg_names.append(a.expr.name)
+        return FunctionDecl(
+            FunctionSignature(
+                arg_types=tuple(arg_types),
+                arg_names=tuple(arg_names),
+                arg_defaults=(None,) * len(self.args),
+                return_type=self.res.tp.to_var(),
+            ),
         )
 
+    @property
+    def egg_name(self) -> None | str:
+        return None
+
 
 @dataclass(frozen=True)
 class FunctionRef:
-    name: str | UnnamedFunctionRef
+    name: str
 
 
 @dataclass(frozen=True)
@@ -380,7 +396,14 @@ class PropertyRef:
 
 
 CallableRef: TypeAlias = (
-    FunctionRef | ConstantRef | MethodRef | ClassMethodRef | InitRef | ClassVariableRef | PropertyRef
+    FunctionRef
+    | ConstantRef
+    | MethodRef
+    | ClassMethodRef
+    | InitRef
+    | ClassVariableRef
+    | PropertyRef
+    | UnnamedFunctionRef
 )
 
 

python/egglog/egraph.py

@@ -694,7 +694,7 @@ def _fn_decl(
     is_builtin: bool,
     ruleset: Ruleset | None = None,
     unextractable: bool = False,
-) -> tuple[FunctionRef | MethodRef | PropertyRef | ClassMethodRef | InitRef, Callable[[], None]]:
+) -> tuple[CallableRef, Callable[[], None]]:
     """
     Sets the function decl for the function object and returns the ref as well as a thunk that sets the default callable.
     """
@@ -707,10 +707,7 @@ def _fn_decl(
     if "Callable" not in hint_globals:
         hint_globals["Callable"] = Callable
 
-    try:
-        hints = get_type_hints(fn, hint_globals, hint_locals)
-    except Exception as e:
-        raise TypeError(f"Failed to get type hints for {fn}") from e
+    hints = get_type_hints(fn, hint_globals, hint_locals)
 
     params = list(signature(fn).parameters.values())
 
@@ -771,33 +768,41 @@ def _fn_decl(
         )
     )
     decls.update(*merge_action)
-
-    signature_ = FunctionSignature(
-        return_type=None if mutates_first_arg else return_type,
-        var_arg_type=var_arg_type,
-        arg_types=arg_types,
-        arg_names=arg_names,
-        arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
-    )
-    decl = FunctionDecl(
-        signature=signature_,
-        cost=cost,
-        egg_name=egg_name,
-        merge=merged.__egg_typed_expr__.expr if merged is not None else None,
-        unextractable=unextractable,
-        builtin=is_builtin,
-        default=None if default is None else default.__egg_typed_expr__.expr,
-        on_merge=tuple(a.action for a in merge_action),
-    )
-    res = Thunk.fn(_create_default_value, decls, ref, fn, signature_, ruleset)
+    # defer this in generator so it doesnt resolve for builtins eagerly
+    args = (TypedExprDecl(tp.to_just(), VarDecl(name, False)) for name, tp in zip(arg_names, arg_types, strict=True))
+    res_ref: FunctionRef | MethodRef | ClassMethodRef | PropertyRef | InitRef | UnnamedFunctionRef
+    res_thunk: Callable[[], object]
     # If we were not passed in a ref, this is an unnamed funciton, so eagerly compute the value and use that to refer to it
     if not ref:
-        res_value = res()
-        assert isinstance(res_value, RuntimeExpr)
-        just_arg_types = tuple(tp.to_just() for tp in arg_types)
-        ref = FunctionRef(UnnamedFunctionRef(just_arg_types, arg_names, res_value.__egg_typed_expr__))
-    decls.set_function_decl(ref, decl)
-    return ref, Thunk.fn(_add_default_rewrite_function, decls, ref, signature_.semantic_return_type, ruleset, res)
+        tuple_args = tuple(args)
+        res = _create_default_value(decls, ref, fn, tuple_args, ruleset)
+        assert isinstance(res, RuntimeExpr)
+        res_ref = UnnamedFunctionRef(tuple_args, res.__egg_typed_expr__)
+        decls._unnamed_functions.add(res_ref)
+        res_thunk = Thunk.value(res)
+
+    else:
+        signature_ = FunctionSignature(
+            return_type=None if mutates_first_arg else return_type,
+            var_arg_type=var_arg_type,
+            arg_types=arg_types,
+            arg_names=arg_names,
+            arg_defaults=tuple(a.__egg_typed_expr__.expr if a is not None else None for a in arg_defaults),
+        )
+        decl = FunctionDecl(
+            signature=signature_,
+            cost=cost,
+            egg_name=egg_name,
+            merge=merged.__egg_typed_expr__.expr if merged is not None else None,
+            unextractable=unextractable,
+            builtin=is_builtin,
+            default=None if default is None else default.__egg_typed_expr__.expr,
+            on_merge=tuple(a.action for a in merge_action),
+        )
+        res_ref = ref
+        decls.set_function_decl(ref, decl)
+        res_thunk = Thunk.fn(_create_default_value, decls, ref, fn, args, ruleset)
+    return res_ref, Thunk.fn(_add_default_rewrite_function, decls, res_ref, return_type, ruleset, res_thunk)
 
 
 # Overload to support aritys 0-4 until variadic generic support map, so we can map from type to value
@@ -872,38 +877,25 @@ def _create_default_value(
     decls: Declarations,
     ref: CallableRef | None,
     fn: Callable,
-    signature: FunctionSignature,
+    args: Iterable[TypedExprDecl],
     ruleset: Ruleset | None,
 ) -> object:
-    args: list[object] = [
-        RuntimeExpr.__from_values__(
-            decls,
-            TypedExprDecl(
-                tp.to_just(),
-                VarDecl(name, False),
-            ),
-        )
-        for name, tp in zip(signature.arg_names, signature.arg_types, strict=False)
-    ]
+    args: list[object] = [RuntimeExpr.__from_values__(decls, a) for a in args]
 
     # If this is a classmethod, add the class as the first arg
     if isinstance(ref, ClassMethodRef):
         tp = decls.get_paramaterized_class(ref.class_name)
         args.insert(0, RuntimeClass(Thunk.value(decls), tp))
     with set_current_ruleset(ruleset):
-        try:
-            return fn(*args)
-        except Exception as err:
-            msg = f"Error when calling {fn}"
-            raise ValueError(msg) from err
+        return fn(*args)
 
 
 def _add_default_rewrite_function(
     decls: Declarations,
     ref: CallableRef,
     res_type: TypeOrVarRef,
     ruleset: Ruleset | None,
-    value_thunk: Thunk,
+    value_thunk: Callable[[], object],
 ) -> None:
     """
     Helper functions that resolves a value thunk to create the default value.

python/egglog/egraph_state.py

@@ -401,16 +401,8 @@ def _generate_callable_egg_name(self, ref: CallableRef) -> str:
         """
         match ref:
             case FunctionRef(name):
-                match name:
-                    case str(name):
-                        return name
-                    case UnnamedFunctionRef(arg_types, arg_names, val):
-                        parts = (
-                            list(arg_names)
-                            + [self.type_ref_to_egg(tp) for tp in arg_types]
-                            + [str(self.typed_expr_to_egg(val, False))]
-                        )
-                        return "_".join(parts)
+                return name
+
             case ConstantRef(name):
                 return name
             case (
@@ -422,6 +414,11 @@ def _generate_callable_egg_name(self, ref: CallableRef) -> str:
                 return f"{cls_name}.{name}"
             case InitRef(cls_name):
                 return f"{cls_name}.__init__"
+            case UnnamedFunctionRef(args, val):
+                parts = [str(self._expr_to_egg(a.expr)) + "-" + str(self.type_ref_to_egg(a.tp)) for a in args] + [
+                    str(self.typed_expr_to_egg(val, False))
+                ]
+                return "_".join(parts)
             case _:
                 assert_never(ref)
 

python/egglog/pretty.py

@@ -179,7 +179,7 @@ def __call__(self, decl: AllDecls, toplevel: bool = False) -> None:  # noqa: C90
                     self(de)
             case CallDecl(ref, exprs, _):
                 match ref:
-                    case FunctionRef(UnnamedFunctionRef(_, _, res)):
+                    case FunctionRef(UnnamedFunctionRef(_, res)):
                         self(res.expr)
                     case _:
                         for e in exprs:
@@ -427,6 +427,8 @@ def _call_inner(  # noqa: PLR0911
             case InitRef(class_name):
                 tp_ref = JustTypeRef(class_name, bound_tp_params or ())
                 return str(tp_ref), args
+            case UnnamedFunctionRef():
+                return ref, args
         assert_never(ref)
 
     def _generate_name(self, typ: str) -> str:
@@ -453,9 +455,9 @@ def _pretty_partial(self, ref: CallableRef, args: list[ExprDecl]) -> str:
         """
         match ref:
             case FunctionRef(name):
-                if not isinstance(name, str):
-                    return self._pretty_function_body(name, args)
                 fn = name
+            case UnnamedFunctionRef():
+                return self._pretty_function_body(ref, args)
             case (
                 ClassMethodRef(class_name, method_name)
                 | MethodRef(class_name, method_name)
@@ -484,16 +486,14 @@ def _pretty_function_body(self, fn: UnnamedFunctionRef, args: list[ExprDecl]) ->
         """
         Pretty print the body of a function, partially applying some arguments.
         """
-        var_args = [
-            TypedExprDecl(tp, VarDecl(name, False)) for tp, name in zip(fn.arg_types, fn.arg_names, strict=True)
-        ]
+        var_args = fn.args
         replacements = {var_arg: TypedExprDecl(var_arg.tp, arg) for var_arg, arg in zip(var_args, args, strict=False)}
         var_args = var_args[len(args) :]
         res = replace_typed_expr(fn.res, replacements)
-        arg_names = fn.arg_names[len(args) :]
+        arg_names = fn.args[len(args) :]
         prefix = "lambda"
         if arg_names:
-            prefix += f" {', '.join(arg_names)}"
+            prefix += f" {', '.join(self(a.expr) for a in arg_names)}"
         return f"{prefix}: {self(res.expr)}"