[red-knot] simplify subtypes from unions

[carljm]

Add Type::is_subtype_of method, and simplify subtypes out of unions.

[github-actions]

ruff-ecosystem results

Linter (stable)

✅ ecosystem check detected no linter changes.

Linter (preview)

✅ ecosystem check detected no linter changes.

[MichaReiser]

This change makes UnionBuilder::map O(N^2) but I don't think there's a way to avoid that

[carljm]

This change makes UnionBuilder::map O(N^2) but I don't think there's a way to avoid that

I assume you mean UnionBuilder::add? There is no UnionBuilder::map.

But yes, I agree; it's now O(n^2), and I don't think that's avoidable.

EDIT: oh, I'm guessing you actually meant UnionType::map. UnionBuilder::add is not really O(n^2); it is O(n*m) if you add a union to the builder.

[AlexWaygood]

If we switch to using a Vec for the elements of a union, that will make tests such as union.contains(Type::Any) O(n), right? Not sure how important that is to consider

[carljm]

If we switch to using a Vec for the elements of a union, that will make tests such as union.contains(Type::Any) O(n), right? Not sure how important that is to consider

I don't think in general this is an operation we will need often, for the same reason -- the question will usually be about subtyping or assignability (or equivalence), not simple type equality.

It's possible we will have the specific case of needing to know if Any/Unknown is in the union, but I think if that's an issue we could store an extra boolean flag on every union (and potentially even leave the actual Any/Unknown entry out of it) for less cost than the cost of the FxOrderSet.

[MichaReiser]

EDIT: oh, I'm guessing you actually meant UnionType::map. UnionBuilder::add is not really O(n^2); it is O(n*m) if you add a union to the builder.

Yes sorry. I still think it is because we loop over n elements and loops over all elements that have been added to this point. So it's probably n(n-1)/2

[carljm]

I still think it is because we loop over n elements and loops over all elements that have been added to this point. So it's probably n(n-1)/2

Yes, I agree that this is accurate for UnionType::map.

I was correcting myself about UnionBuilder::add, which is itself only O(n), or O(n*m) if we are adding another union (which won't ever be the case from UnionType::map.)

[hauntsaninja]

In equivalent mypy code, I had to add a special fast path for literals. You can do better than quadratic for unions with lots of literals of the same type, which turns out to be a thing in the wild

[carljm]

Interesting, thanks for pointing this out! I took a look at that optimization in mypy.

I think the case this would optimize is where a union already contains e.g. str, and then we try to add lots of string literal types to it, every one of which is redundant because its a subtype of str. Rather than going through all existing union members to check if each literal is a subtype of any of them, we can keep a hash-set of "types present in this union which have literal forms" and do an O(1) contains check against that set as the first step when adding a literal type to the union. Framed in more general terms, it's identifying that a certain set of common types have a single super-type that is most likely to rule them out of the union, and so we optimize checking for that most likely super-type by identity.

This makes sense; I'd prefer to wait to add this kind of optimization until we see it crop up in a real-world codebase and can evaluate the actual impact of the optimization in our case, but it's definitely a useful idea to keep in mind.

[AlexWaygood]

I think what would be useful is if we added one (or more) benchmarks based on a real-world codebase that makes heavy use of large literals. (I.e., pydantic.)

[hauntsaninja]

That's not quite the right description, it's also useful when the union doesn't contain the supertype (i.e. str). For instance, say if you were combining two unions that you knew consisted only of literal types, you could use a set union, which is linear. The mypy optimisation I added is basically that, but also works when there are non-literal types thrown in as well. Fair enough on waiting though!

[carljm]

Oh, thanks, yeah, I misread the code. The set is unduplicated_literal_fallbacks, not duplicated_literal_fallbacks. So it looks like it's optimizing only the case you described; the mirror image of the case I described.

[AlexWaygood]

I've created a new issue collating some of the perf issues mypy and pyright have encountered relating to unions: #13549

[hauntsaninja]

Nice!
I think there are some mypy PRs missing from the list, so if you're interested in code I'd make sure to look at main.
I'll also make it such that if you're interested in real world use cases you should only have to look at primer, looks like there are 1-2 things I never actually added.