Skip to content

[ty] Improve union builder performance #22048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MichaReiser merged 4 commits into from
Dec 19, 2025
Merged

[ty] Improve union builder performance #22048

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
268a8f8
[ty] Small union builder nits
MichaReiser Dec 18, 2025
0a937d2
Defer insertion
MichaReiser Dec 18, 2025
368405a
More lazy negated computations
MichaReiser Dec 18, 2025
3848ada
Code review feedback, revert in-place removal changes
MichaReiser Dec 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 33 additions & 22 deletions crates/ty_python_semantic/src/types/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,7 +365,7 @@ impl<'db> UnionBuilder<'db> {
Type::StringLiteral(literal) => {
let mut found = None;
let mut to_remove = None;
let ty_negated = ty.negate(self.db);
let mut ty_negated = None;
Copy link
Member

@AlexWaygood AlexWaygood Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

did you consider using a OnceCell here? It might make the logic more readable and less repetitive

Copy link
Member Author

@MichaReiser MichaReiser Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure what the benefit of using OnceCell here is. The main advantage of OnceCell is that it doesn't require mut, but the mut isn't an issue here.

Copy link
Member

@AlexWaygood AlexWaygood Dec 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, yeah, a OnceCell doesn't really make much sense here. The main thing I was wondering about was whether there was any way to reduce the repetition of having to do ty_negated.get_or_insert_with(|| ty.negate(db)) so many times. See #22082.

for (index, element) in self.elements.iter_mut().enumerate() {
match element {
UnionElement::StringLiterals(literals) => {
Expand All @@ -378,13 +378,17 @@ impl<'db> UnionBuilder<'db> {
continue;
}
UnionElement::Type(existing) => {
// e.g. `existing` could be `Literal[""] & Any`,
// and `ty` could be `Literal[""]`
if ty.is_subtype_of(self.db, *existing) {
return;
}
if existing.is_subtype_of(self.db, ty) {
Copy link
Member

@AlexWaygood AlexWaygood Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had to play around with this code a little to figure out how we even get to this branch, and why it's correct here for to_remove to be Option<usize> rather than Vec<usize>. It might be helpful to add a comment here (and similar comments to the if existing.is_subtype_of(db, ty) calls in the Type::IntLiteral() and Type::BytesLiteral branches below):

Suggested change
if existing.is_subtype_of(self.db, ty) {
// e.g. `existing` could be `Literal[""] & Any`,
// and `ty` could be `Literal[""]`
if existing.is_subtype_of(self.db, ty) {

Copy link
Contributor

@carljm carljm Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes I agree those comments would be helpful.

I think the reason it's OK to just track one to_remove is also a bit subtle -- it's because there are a limited number of possible subtypes of a literal, and all the possible subtypes (e.g. Literal[1] & Any, Literal[1] & Unknown) are also redundant with each other, so it's not possible that we'd have more than one as an existing element.

to_remove = Some(index);
continue;
Copy link
Member Author

@MichaReiser MichaReiser Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should continue here, the same as in push_type?

Copy link
Member

@AlexWaygood AlexWaygood Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, I think this is correct

}
if ty_negated.is_subtype_of(self.db, *existing) {
let negated = ty_negated.get_or_insert_with(|| ty.negate(self.db));
if negated.is_subtype_of(self.db, *existing) {
// The type that includes both this new element, and its negation
// (or a supertype of its negation), must be simply `object`.
self.collapse_to_object();
Expand All @@ -410,7 +414,7 @@ impl<'db> UnionBuilder<'db> {
Type::BytesLiteral(literal) => {
let mut found = None;
let mut to_remove = None;
let ty_negated = ty.negate(self.db);
let mut ty_negated = None;
for (index, element) in self.elements.iter_mut().enumerate() {
match element {
UnionElement::BytesLiterals(literals) => {
Expand All @@ -428,8 +432,11 @@ impl<'db> UnionBuilder<'db> {
}
if existing.is_subtype_of(self.db, ty) {
to_remove = Some(index);
continue;
}
if ty_negated.is_subtype_of(self.db, *existing) {

let negated = ty_negated.get_or_insert_with(|| ty.negate(self.db));
if negated.is_subtype_of(self.db, *existing) {
// The type that includes both this new element, and its negation
// (or a supertype of its negation), must be simply `object`.
self.collapse_to_object();
Expand All @@ -455,7 +462,7 @@ impl<'db> UnionBuilder<'db> {
Type::IntLiteral(literal) => {
let mut found = None;
let mut to_remove = None;
let ty_negated = ty.negate(self.db);
let mut ty_negated = None;
for (index, element) in self.elements.iter_mut().enumerate() {
match element {
UnionElement::IntLiterals(literals) => {
Expand All @@ -473,8 +480,11 @@ impl<'db> UnionBuilder<'db> {
}
if existing.is_subtype_of(self.db, ty) {
to_remove = Some(index);
continue;
}
if ty_negated.is_subtype_of(self.db, *existing) {

let negated = ty_negated.get_or_insert_with(|| ty.negate(self.db));
if negated.is_subtype_of(self.db, *existing) {
// The type that includes both this new element, and its negation
// (or a supertype of its negation), must be simply `object`.
self.collapse_to_object();
Expand Down Expand Up @@ -549,18 +559,14 @@ impl<'db> UnionBuilder<'db> {
// unpacking them.
let should_simplify_full = !matches!(ty, Type::TypeAlias(_)) && !self.cycle_recovery;

let mut ty_negated: Option<Type> = None;
let mut to_remove = SmallVec::<[usize; 2]>::new();
let ty_negated = if should_simplify_full {
ty.negate(self.db)
} else {
Type::Never // won't be used
};

for (index, element) in self.elements.iter_mut().enumerate() {
for (i, element) in self.elements.iter_mut().enumerate() {
let element_type = match element.try_reduce(self.db, ty) {
ReduceResult::KeepIf(keep) => {
if !keep {
to_remove.push(index);
to_remove.push(i);
}
continue;
}
Expand All @@ -587,19 +593,22 @@ impl<'db> UnionBuilder<'db> {
// problematic if some of those fields point to recursive `Union`s. To avoid cycles,
// compare `TypedDict`s by name/identity instead of using the `has_relation_to`
// machinery.
if let (Type::TypedDict(element_td), Type::TypedDict(ty_td)) = (element_type, ty) {
if element_td == ty_td {
return;
}
Comment on lines -591 to -593
Copy link
Member Author

@MichaReiser MichaReiser Dec 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This check is redundant with the check on line 577

if element_type.is_typed_dict() && ty.is_typed_dict() {
continue;
}

if should_simplify_full && !matches!(element_type, Type::TypeAlias(_)) {
if ty.is_redundant_with(self.db, element_type) {
return;
} else if element_type.is_redundant_with(self.db, ty) {
to_remove.push(index);
} else if ty_negated.is_subtype_of(self.db, element_type) {
}

if element_type.is_redundant_with(self.db, ty) {
to_remove.push(i);
continue;
}

let negated = ty_negated.get_or_insert_with(|| ty.negate(self.db));
if negated.is_subtype_of(self.db, element_type) {
// We add `ty` to the union. We just checked that `~ty` is a subtype of an
// existing `element`. This also means that `~ty | ty` is a subtype of
// `element | ty`, because both elements in the first union are subtypes of
Expand All @@ -614,10 +623,12 @@ impl<'db> UnionBuilder<'db> {
}
}
}
if let Some((&first, rest)) = to_remove.split_first() {

let mut to_remove = to_remove.into_iter();
if let Some(first) = to_remove.next() {
self.elements[first] = UnionElement::Type(ty);
// We iterate in descending order to keep remaining indices valid after `swap_remove`.
for &index in rest.iter().rev() {
for index in to_remove.rev() {
self.elements.swap_remove(index);
}
} else {
Expand Down
Loading