egglog 7.0.0__cp311-none-win_amd64.whl → 7.2.0__cp311-none-win_amd64.whl

egglog/egraph_state.py

@@ -19,7 +19,7 @@ from .type_constraint_solver import TypeConstraintError, TypeConstraintSolver
 if TYPE_CHECKING:
     from collections.abc import Iterable
 
-__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT"]
+__all__ = ["EGraphState", "GLOBAL_PY_OBJECT_SORT", "_rule_var_name"]
 
 # Create a global sort for python objects, so we can store them without an e-graph instance
 # Needed when serializing commands to egg commands when creating modules
@@ -87,17 +87,27 @@ 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
+                    cmd = self.command_to_egg(rule, name)
+                    self.egraph.run_program(cmd)
+                    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:
@@ -106,8 +116,8 @@ class EGraphState:
             case RewriteDecl(tp, lhs, rhs, conditions) | BiRewriteDecl(tp, lhs, rhs, conditions):
                 self.type_ref_to_egg(tp)
                 rewrite = bindings.Rewrite(
-                    self.expr_to_egg(lhs),
-                    self.expr_to_egg(rhs),
+                    self._expr_to_egg(lhs),
+                    self._expr_to_egg(rhs),
                     [self.fact_to_egg(c) for c in conditions],
                 )
                 return (
@@ -121,6 +131,16 @@ class EGraphState:
                     [self.fact_to_egg(f) for f in body],
                 )
                 return bindings.RuleCommand(name or "", ruleset, rule)
+            case DefaultRewriteDecl(ref, expr):
+                decl = self.__egg_decls__.get_callable_decl(ref).to_function_decl()
+                sig = decl.signature
+                assert isinstance(sig, FunctionSignature)
+                args = tuple(
+                    TypedExprDecl(tp.to_just(), VarDecl(_rule_var_name(name)))
+                    for name, tp in zip(sig.arg_names, sig.arg_types, strict=False)
+                )
+                rewrite_decl = RewriteDecl(sig.semantic_return_type.to_just(), CallDecl(ref, args), expr, (), False)
+                return self.command_to_egg(rewrite_decl, ruleset)
             case _:
                 assert_never(cmd)
 
@@ -130,13 +150,13 @@ class EGraphState:
                 return bindings.Let(name, self.typed_expr_to_egg(typed_expr))
             case SetDecl(tp, call, rhs):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
-                return bindings.Set(call_.name, call_.args, self.expr_to_egg(rhs))
+                call_ = self._expr_to_egg(call)
+                return bindings.Set(call_.name, call_.args, self._expr_to_egg(rhs))
             case ExprActionDecl(typed_expr):
                 return bindings.Expr_(self.typed_expr_to_egg(typed_expr))
             case ChangeDecl(tp, call, change):
                 self.type_ref_to_egg(tp)
-                call_ = self.expr_to_egg(call)
+                call_ = self._expr_to_egg(call)
                 egg_change: bindings._Change
                 match change:
                     case "delete":
@@ -148,7 +168,7 @@ class EGraphState:
                 return bindings.Change(egg_change, call_.name, call_.args)
             case UnionDecl(tp, lhs, rhs):
                 self.type_ref_to_egg(tp)
-                return bindings.Union(self.expr_to_egg(lhs), self.expr_to_egg(rhs))
+                return bindings.Union(self._expr_to_egg(lhs), self._expr_to_egg(rhs))
             case PanicDecl(name):
                 return bindings.Panic(name)
             case _:
@@ -158,7 +178,7 @@ class EGraphState:
         match fact:
             case EqDecl(tp, exprs):
                 self.type_ref_to_egg(tp)
-                return bindings.Eq([self.expr_to_egg(e) for e in exprs])
+                return bindings.Eq([self._expr_to_egg(e) for e in exprs])
             case ExprFactDecl(typed_expr):
                 return bindings.Fact(self.typed_expr_to_egg(typed_expr))
             case _:
@@ -184,14 +204,16 @@ 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,
+                        self._expr_to_egg(decl.default) if decl.default else None,
+                        self._expr_to_egg(decl.merge) if decl.merge else None,
                         [self.action_to_egg(a) for a in decl.on_merge],
                         decl.cost,
                         decl.unextractable,
@@ -212,25 +234,30 @@ 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.
         if decl.builtin:
             for method in decl.class_methods:
                 self.callable_ref_to_egg(ClassMethodRef(ref.name, method))
+            if decl.init:
+                self.callable_ref_to_egg(InitRef(ref.name))
 
         return egg_name
 
@@ -247,15 +274,15 @@ class EGraphState:
 
     def typed_expr_to_egg(self, typed_expr_decl: TypedExprDecl) -> bindings._Expr:
         self.type_ref_to_egg(typed_expr_decl.tp)
-        return self.expr_to_egg(typed_expr_decl.expr)
+        return self._expr_to_egg(typed_expr_decl.expr)
 
     @overload
-    def expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
+    def _expr_to_egg(self, expr_decl: CallDecl) -> bindings.Call: ...
 
     @overload
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr: ...
 
-    def expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
+    def _expr_to_egg(self, expr_decl: ExprDecl) -> bindings._Expr:
         """
         Convert an ExprDecl to an egg expression.
 
@@ -292,6 +319,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)
 
@@ -338,6 +368,8 @@ def _generate_callable_egg_name(ref: CallableRef) -> str:
             | PropertyRef(cls_name, name)
         ):
             return f"{cls_name}_{name}"
+        case InitRef(cls_name):
+            return f"{cls_name}___init__"
         case _:
             assert_never(ref)
 
@@ -371,26 +403,50 @@ 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()
-            if isinstance(callable_ref, ClassMethodRef):
-                possible_types = self.state._get_possible_types(callable_ref.class_name)
+            signature = self.decls.get_callable_decl(callable_ref).to_function_decl().signature
+            assert isinstance(signature, FunctionSignature)
+            if isinstance(callable_ref, ClassMethodRef | InitRef | MethodRef):
+                # Need OR in case we have class method whose class whas never added as a sort, which would happen
+                # if the class method didn't return that type and no other function did. In this case, we don't need
+                # to care about the type vars and we we don't need to bind any possible type.
+                possible_types = self.state._get_possible_types(callable_ref.class_name) or [None]
                 cls_name = callable_ref.class_name
             else:
                 possible_types = [None]
@@ -399,19 +455,33 @@ class FromEggState:
                 tcs = TypeConstraintSolver(self.decls)
                 if possible_type and possible_type.args:
                     tcs.bind_class(possible_type)
-
                 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,
+                    # Don't include bound type params if this is just a method, we only needed them for type resolution
+                    # but dont need to store them
+                    bound_tp_params if isinstance(callable_ref, ClassMethodRef | InitRef) else None,
+                )
         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
+
+
+def _rule_var_name(s: str) -> str:
+    """
+    Create a hidden variable name, for rewrites, so that let bindings or function won't conflict with it
+    """
+    return f"__var__{s}"

egglog/examples/higher_order_functions.py

@@ -0,0 +1,50 @@
+# mypy: disable-error-code="empty-body"
+"""
+Higher Order Functions
+======================
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from egglog import *
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Math(Expr):
+    def __init__(self, i: i64Like) -> None: ...
+
+    def __add__(self, other: Math) -> Math: ...
+
+
+class MathList(Expr):
+    def __init__(self) -> None: ...
+
+    def append(self, i: Math) -> MathList: ...
+
+    def map(self, f: Callable[[Math], Math]) -> MathList: ...
+
+
+@ruleset
+def math_ruleset(i: i64, j: i64, xs: MathList, x: Math, f: Callable[[Math], Math]):  # noqa: ANN201
+    yield rewrite(Math(i) + Math(j)).to(Math(i + j))
+    yield rewrite(xs.append(x).map(f)).to(xs.map(f).append(f(x)))
+    yield rewrite(MathList().map(f)).to(MathList())
+
+
+@function(ruleset=math_ruleset)
+def increment_by_one(x: Math) -> Math:
+    return x + Math(1)
+
+
+egraph = EGraph()
+x = egraph.let("x", MathList().append(Math(1)).append(Math(2)))
+y = egraph.let("y", x.map(increment_by_one))
+egraph.run(math_ruleset.saturate())
+
+egraph.check(eq(y).to(MathList().append(Math(2)).append(Math(3))))
+
+egraph

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)),
     ]
 
 
@@ -882,7 +882,10 @@ converter(IntOrTuple, OptionalIntOrTuple, OptionalIntOrTuple.some)
 
 @function
 def asarray(
-    a: NDArray, dtype: OptionalDType = OptionalDType.none, copy: OptionalBool = OptionalBool.none
+    a: NDArray,
+    dtype: OptionalDType = OptionalDType.none,
+    copy: OptionalBool = OptionalBool.none,
+    device: OptionalDevice = OptionalDevice.none,
 ) -> NDArray: ...
 
 
@@ -1346,7 +1349,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:
     """
@@ -162,6 +166,8 @@ class TraverseContext:
                 pass
             case EqDecl(_, decls) | SequenceDecl(decls) | RulesetDecl(decls):
                 for de in decls:
+                    if isinstance(de, DefaultRewriteDecl):
+                        continue
                     self(de)
             case CallDecl(_, exprs, _):
                 for e in exprs:
@@ -170,6 +176,12 @@ class TraverseContext:
                 if until:
                     for f in until:
                         self(f)
+            case PartialCallDecl(c):
+                self(c)
+            case CombinedRulesetDecl(_):
+                pass
+            case DefaultRewriteDecl():
+                pass
             case _:
                 assert_never(decl)
 
@@ -231,6 +243,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):
@@ -265,8 +280,12 @@ class PrettyContext:
             case RulesetDecl(rules):
                 if ruleset_name:
                     return f"ruleset(name={ruleset_name!r})", f"ruleset_{ruleset_name}"
-                args = ", ".join(map(self, rules))
+                args = ", ".join(self(r) for r in rules if not isinstance(r, DefaultRewriteDecl))
                 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):
@@ -283,6 +302,9 @@ class PrettyContext:
                     return ruleset_str, "schedule"
                 args = ", ".join(map(self, until))
                 return f"run({ruleset_str}, {args})", "schedule"
+            case DefaultRewriteDecl():
+                msg = "default rewrites should not be pretty printed"
+                raise TypeError(msg)
         assert_never(decl)
 
     def _call(
@@ -302,19 +324,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)
@@ -346,10 +377,8 @@ class PrettyContext:
             case FunctionRef(name):
                 return name, args
             case ClassMethodRef(class_name, method_name):
-                fn_str = str(JustTypeRef(class_name, bound_tp_params or ()))
-                if method_name != "__init__":
-                    fn_str += f".{method_name}"
-                return fn_str, args
+                tp_ref = JustTypeRef(class_name, bound_tp_params or ())
+                return f"{tp_ref}.{method_name}", args
             case MethodRef(_class_name, method_name):
                 slf, *args = args
                 slf = self(slf, parens=True)
@@ -376,6 +405,9 @@ class PrettyContext:
                 return f"{class_name}.{variable_name}"
             case PropertyRef(_class_name, property_name):
                 return f"{self(args[0], parens=True)}.{property_name}"
+            case InitRef(class_name):
+                tp_ref = JustTypeRef(class_name, bound_tp_params or ())
+                return str(tp_ref), args
         assert_never(ref)
 
     def _generate_name(self, typ: str) -> str:
@@ -397,6 +429,30 @@ 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 InitRef(class_name):
+            return class_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