RFC: Higher-Order Function Calls

[rchen152]

What is this?

A proposed redesign of how Pyrefly models calls to higher-order functions.

Why?

Currently, Pyrefly tracks correlations between type parameters introduced by higher-order function calls through sets of variables that are individually solved to Type::CallableResidual objects, which are then used to reconstruct generic or overloaded types. This scheme is powerful and expressive, handling many cases that other type checkers cannot, but also terrifically complex, producing a slow drip of hard-to-fix bugs.

This proposal presents an alternative design in which most of the complexity is localized to short-lived data structures in the variable solver.

Overview

The design has two halves:

1. Marking and finalizing free type parameters in generic functions.
2. Building and solving tables of correlations for overloaded functions.

The first half is an incremental improvement to how things work today; the second half is a ground-up redesign.

Generic functions

Consider this example of a generic function (identity) passed to a higher-order function (apply):

from typing import Callable

def apply[P, R](f: Callable[[P], R]) -> Callable[[P], R]: ...
def identity[X](x: X) -> X: ...

reveal_type(apply(identity))  # [R](R) -> R

To determine the type returned from apply(identity), Pyrefly checks the type of the identity argument against the expected type:

[X](X) -> X <: [P, R](P) -> R

In the process, we create fresh variables - @X, @P, @R - for the type parameters involved, then unify the variables. The order in which unification happens doesn’t matter too much, but let’s say @X and @P are solved to @R, so at the end of the check, we have one unsolved variable left, @R.

Today, the variable is solved to a residual type, GenericResidual@R. A finalization method runs on the returned type, converting every occurrence of GenericResidual@R to a type parameter R and making the callable generic over R.

Proposed change: solve the variable directly to a type parameter and set a new free flag on the type parameter. The finalization method finds the innermost callable whose parameters reference a free type parameter, makes the callable generic over it, and unsets free.

This gains us:

• Simplicity: using a flag on type parameters allows us to delete a chunk of the plumbing for Type::CallableResidual.
• Correctness: in a type with nested callables, the current finalization strategy sometimes makes the wrong callable generic.

Aside: why a permanent free flag on a type rather than transient solver data? There are cases in which the type parameter needs to remain free in the returned type:

from typing import Callable

class Wrapper[**P, R]:
    fn: Callable[P, R]

def wrap[P, R](f: Callable[[P], R]) -> Wrapper[[P], R]: ...
def identity[X](x: X) -> X: ...

# `Wrapper[[R], R]` - R is free
wrapped = reveal_type(wrap(identity))
# `[R](R) -> R` - when we access `fn`, we see that `R` is free and make `fn` generic over `R`
reveal_type(wrapped.fn)

Overloaded functions

Consider this example of an overloaded function (parse) passed to a higher-order function (apply):

from typing import Callable, overload

@overload
def parse(x: int) -> str: ...
@overload
def parse(x: str) -> int: ...
def parse(x: int | str) -> str | int: ...

def apply[A, B](f: Callable[[A], B]) -> Callable[[A], B]: ...

reveal_type(apply(parse))  # Overload[(int) -> str, (str) -> int]

I won’t go into detail about how Pyrefly evaluates apply(parse) today, but the key observation is that it uses overload residuals, stand-ins for individual type parameters that record what each signature of the overload would solve that type parameter to. However, unlike generic residuals - for which we get correlation for free via variable unification - overload residuals essentially end up being free-floating function parts, and every residual has to carry enough data to reconstruct how it correlates with other residuals.

I propose to instead have the solver build a table of correlations, in which the rows are overload signatures and the columns are type parameters. We can eliminate rows that solve type parameters to conflicting types, read off the returned types from the surviving rows, and combine the types for our final answer.

Here’s what the table would look like for apply(parse) above:

	A	B
(int) -> str	int	str
(str) -> int	str	int

There are no conflicting solutions, so we keep both rows. From the first row, we read off A=int, B=str, giving a returned type of (int) -> str. From the second row, we read off A=str, B=int, giving a returned type of (str) -> int. Combine them, and we get our desired answer of Overload[(int) -> str, (str) -> int].

Here are a couple more involved examples showing when we would eliminate rows. This example demonstrates pruning:

from typing import Callable, overload

@overload
def parse(x: int) -> str: ...
@overload
def parse(x: str) -> int: ...
def parse(x: int | str) -> str | int: ...

def pick[A, B](f: Callable[[A], B], a: A) -> B: ...

reveal_type(pick(parse, a=1))   # str

	A	B
(int) -> str	int	str
(str) -> int	str	int

Passing in a=1 solves type parameter A to int, so we prune the row with A=str, leaving exactly one row that gives a returned type of str.

Pruning has two interesting edge cases:

• If pruning eliminates all rows, we report an incompatible-overload-argument error, renamed from incompatible-overload-residual because residuals no longer exist as a user-facing concept.
• If pruning cannot eliminate rows because a type parameter is solved to a gradual type, we use the same rule that overload evaluation does in this case: fall back to the “most general” type if one exists and otherwise to Any.

This example demonstrates joining:

from typing import Callable, overload

@overload
def parse(x: int) -> str: ...
@overload
def parse(x: str) -> int: ...
def parse(x: int | str) -> str | int: ...

@overload
def fmt(x: str) -> bytes: ...
@overload
def fmt(x: int) -> bool: ...
def fmt(x: str | int) -> bytes | bool: ...

def compose[A, B, C](f: Callable[[A], B], g: Callable[[B], C]) -> Callable[[A], C]: ...

reveal_type(compose(parse, fmt))  # Overload[(int) -> bytes, (str) -> bool]

	A	B	C
1. parse: (int) -> str	int	str
2. parse: (str) -> int	str	int
3. fmt: (str) -> bytes		str	bytes
4. fmt: (int) -> bool		int	bool

We join rows 1 and 3 on B=str, and join rows 2 and 4 on B=int, producing the final table:

	A	B	C
parse: (int) -> str, fmt: (str) -> bytes	int	str	bytes
parse: (str) -> int, fmt: (int) -> bool	str	int	bool

From there, we read off the rows and combine them as before.

This design gains us:

• Simplicity: Type::CallableResidual and all of its associated machinery can be deleted. Complexity is (mostly) contained to table operations in the solver. (See Combining rows below for an important caveat.)
• Precision: the join example does not work with residuals, so we currently fall back to an imprecise union.

Combining rows

Above, I intentionally picked examples in which the rows combine into an obvious final type. The table is not always so well-behaved. Consider:

from typing import Callable, overload

@overload
def parse(x: int) -> str: ...
@overload
def parse(x: str) -> int: ...
def parse(x: int | str) -> str | int: ...

class Wrapper[**P, R]:
    def __init__(self, fn: Callable[P, R]):
        self.fn = fn

wrapper = reveal_type(Wrapper(parse))  # ???
wrapper.fn(42)  # legal
wrapper.fn("")  # legal

The per-row returned types are Wrapper[[int], str] and Wrapper[[str], int]. How do we combine them? If we use a union, both of the legal wrapper.fn(...) calls shown above will fail. An intersection isn’t right, because intersections are unordered, whereas our rows are ordered because they are derived from an overloaded function. But we can’t use overloads either, because Wrapper is not a callable.

I propose adding a new type, tentatively called Type::Overloaded, that holds a list of alternatives derived from overload signatures. Semantics:

• Overloaded is assignable to a type if any of its alternatives is. This is consistent with overloaded functions.
• Attribute access looks up the attribute on all alternatives and re-combines the lookup results.
• Calls follow overload semantics.
• Other operations (e.g., subscription) approximate overload semantics: alternatives with incompatible input signatures are dropped, and the results of the operation on the remaining alternatives are re-combined.

In the above example, the type of Wrapper(parse) would be Overloaded[Wrapper[[int], str], Wrapper[[str], int]]. Looking up wrapper.fn would get us Overload[(int) -> str, (str) -> int], and the calls with 42 and "" would succeed as expected.

A nice property of Type::Overloaded is that it can be an input to another correlation table, providing a natural way to preserve overloaded structure through multiple calls.

When should combining rows produce a Type::Overloaded? Proposed rules:

1. If the table has only one surviving row, produce the returned type from that row.
2. Else if the returned types are all callables, produce an overloaded function.
3. Else if the returned types all have the same outer shape, produce an Overloaded.
4. Else, produce a union.

Rule 3 is a heuristic that attempts to cover common cases that can preserve overloaded structure (wrapper classes, callback protocols). It does also catch cases that obviously do not contain overloaded structure (e.g., it will happily build Overloaded[list[int], list[str]]), but the failure mode isn’t too bad - a few false negatives in edge cases. I also considered some more sophisticated heuristics involving things like variance or counting or comparing type arguments, but they were considerably more complicated to implement and still had obvious failure modes.

Finally, the introduction of Type::Overloaded means that this design does not reduce the total number of Type variants in Pyrefly. However, the fiddliest logic is now contained to short-lived correlation tables, and Overloaded’s alternatives are regular types. The only novel thing about it is its overloaded structure, and even that has precedence in how overloaded functions work.

Performance

Joins are expensive. In a prototype of this design, I timed calls to a function that took N copies of an 8-signature overloaded function. Each additional copy made the call about 10x slower. Capping the number of rows in the table is a must.

With a cap of 64 rows, full check times on numpy and scipy-stubs - two overload-heavy projects - were indistinguishable from trunk. When the number of rows exceeds the cap, we use the same gradual fallback rule as if pruning could not eliminate rows.

[grievejia]

Thanks for the detailed writeup! I want to confirm my high-level understanding, then ask two clarifying questions.

High-level model

My reading is that the proposal supports two forms of deferred constraints.

For a generic argument:

def identity[T](x: T) -> T: ...
def apply[A, B](f: Callable[[A], B]) -> Callable[[A], B]: ...

The current solver treats A, T, and B as one equality class:

A = T AND T = B

After removing the argument-local T, one free type parameter represents the
unresolved equality A = B. Substitution into the return type produces
[R](R) -> R.

A free parameter can also survive inside a returned type, such as
Wrapper[[R], R], and become bound later when wrapped.fn is read.

For an overloaded argument:

@overload
def parse(x: int) -> str: ...

@overload
def parse(x: str) -> int: ...

Checking parse against Callable[[A], B] produces a bounded, ordered
disjunction:

(A = int AND B = str)
OR
(A = str AND B = int)

Each disjunct is a table row. A later constraint is distributed into every row:

(row_1 OR row_2) AND C
=
(row_1 AND C) OR (row_2 AND C)

The solver removes unsatisfiable rows and refines rows where C provides a new
solution. Multiple overloaded arguments combine their tables on shared
variables. The surviving rows are substituted into the return type.

Under this reading, free type parameters preserve unresolved equality classes
and substitutions that contain free parameters, while tables preserve finite
ordered disjunctions of assignments.

Is this the intended model?

Table-cell meaning and join compatibility

Do table cells contain exact solved types, lower and upper bounds, or some other
branch-local solver state?

For example:

class Animal: ...
class Dog(Animal): ...
class Cat(Animal): ...

@overload
def make(x: int) -> Dog: ...
@overload
def make(x: str) -> Cat: ...
def make(x: int | str) -> Animal: ...

@overload
def encode(x: Animal) -> bytes: ...
@overload
def encode(x: bytes) -> str: ...
def encode(x: Animal | bytes) -> bytes | str: ...

def compose[A, B, C](
    f: Callable[[A], B],
    g: Callable[[B], C],
) -> Callable[[A], C]: ...

result = compose(make, encode)

The (int) -> Dog row from make and the (Animal) -> bytes row from
encode produce:

Dog <= B
B <= Animal

The conjunction Dog <= B <= Animal is satisfiable. A table that records these
cells as the exact assignments B = Dog and B = Animal rejects the join.
Exact equality rejects every branch pair in this example, so the table becomes
empty and a valid call fails.

The RFC's compose example has B = str on both sides of one join and
B = int on both sides of the other. Those identical cells do not expose this
question.

The solver already tracks a lower and an upper bound for each variable. Is each
table row intended to retain those bounds and run normal bound solving after a
join?

Generic parameters inside overload rows

Can a table row contain free generic parameters or types that contain them?

For example:

@overload
def f(x: int) -> str: ...

@overload
def f[T](x: list[T]) -> T: ...

Checking f against Callable[[A], B] appears to produce:

(A = int AND B = str)
OR
(A = list[free T] AND B = free T)

If a later argument solves A = list[bytes], pruning should bind T = bytes
and remove the concrete row.

The same issue occurs during a join. If another table has B = str, should it
combine with the generic row by binding T = str, producing
A = list[str]?

The current implementation already allows a generic residual inside an
overload branch and finalizes overload and generic residuals in sequence. The
proposed mechanisms must preserve this composition. Where does the row-local
generic state live, and how do pruning and joins update it?

[stroxler]

Thanks for the write-up.

Proposed Design

I think this makes sense. Everything you propose is roughly what I was trying to do with CallableResidual in the first place, but I couldn't make enough sense of the solver to construct these types in the right place at the time so I wound up with a pretty atomized solution.

I think your mental model of the solver is clearer than mine was at the time, plus I think we've also made some improvements since then. If this seems feasible, I'm 100% in support. The Type::Overloaded solution in particular seems like a nice workaround for one of the biggest issues I was struggling with, which was how to defer the overload until method call time.

Thoughts on Jia's response

AND and OR

I think you are sort of right, but I think OR is a slippery concept. Read naively, it would suggest a union. Really I think that for overloads, it's sort of a second flavor of AND (one that implies the intersection-like behavior of overload) that binds looser than the inner AND. I'm not really sure how to map it to simple boolean operators, or if it can be cleanly mapped. The operator maybe should just be called OVERLOAD, if we were going to describe it in terms of a tree of ops on constraints(?)

I don't think this is quite the same thing as true disjunctive constraints, which can happen in some circumstances (although we don't handle them, so we misbehave in those settings).

Table-cell contents

I'm not sure exactly what @rchen152 has planned, but in the existing setup as I remember it:

• Overloads can contain generics, which would suggest that free inside of the overload table probably should be allowed
• At the moment we always fully solve the branches, unless they are unconstrained (in which case they become generic). I think deferring constraints could make sense in principle (for example we could create a free var with just an upper bound in some settings), but I didn't have appetite to add complexity to CallableResidual. Maybe if the table implementation is much simpler we could consider this, but I suspect the payoff is low.