[ty] Skip literal branches in union builder if it contains no literals

crates/ty_python_semantic/resources/mdtest/annotations/literal.md

@@ -304,7 +304,7 @@ def union_example(
         None,
     ],
 ):
-    reveal_type(x)  # revealed: Unknown | Literal[-1, 0, 1, "A", "B", "foo", "bar", b"A", b"\x00", b"\x07", True] | None
+    reveal_type(x)  # revealed: Unknown | Literal[-1, "A", b"A", b"\x00", b"\x07", 0, 1, "B", "foo", "bar", True] | None
 ```
 
 ## Detecting Literal outside typing and typing_extensions

crates/ty_python_semantic/src/types/builder.rs

@@ -232,6 +232,7 @@ pub(crate) struct UnionBuilder<'db> {
     db: &'db dyn Db,
     unpack_aliases: bool,
     order_elements: bool,
+    contains_literals: bool,
     // This is enabled when joining types in a `cycle_recovery` function.
     // Since a cycle cannot be created within a `cycle_recovery` function, execution of `is_redundant_with` is skipped.
     cycle_recovery: bool,
@@ -245,6 +246,7 @@ impl<'db> UnionBuilder<'db> {
             elements: vec![],
             unpack_aliases: true,
             order_elements: false,
+            contains_literals: false,
             cycle_recovery: false,
             recursively_defined: RecursivelyDefined::No,
         }
@@ -362,10 +364,10 @@ impl<'db> UnionBuilder<'db> {
             // add it to, or an existing element that is a super-type of string literals, which
             // means we shouldn't add it. Otherwise, add a new `UnionElement::StringLiterals`
             // containing it.
-            Type::StringLiteral(literal) => {
+            Type::StringLiteral(literal) if self.contains_literals || self.elements.is_empty() => {
                 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::StringLiterals(literals) => {
@@ -383,8 +385,10 @@ 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();
@@ -397,6 +401,7 @@ impl<'db> UnionBuilder<'db> {
                 if let Some(found) = found {
                     found.insert(literal);
                 } else {
+                    self.contains_literals = true;
                     self.elements
                         .push(UnionElement::StringLiterals(FxOrderSet::from_iter([
                             literal,
@@ -407,10 +412,10 @@ impl<'db> UnionBuilder<'db> {
                 }
             }
             // Same for bytes literals as for string literals, above.
-            Type::BytesLiteral(literal) => {
+            Type::BytesLiteral(literal) if self.contains_literals || self.elements.is_empty() => {
                 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) => {
@@ -428,8 +433,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();
@@ -442,6 +450,7 @@ impl<'db> UnionBuilder<'db> {
                 if let Some(found) = found {
                     found.insert(literal);
                 } else {
+                    self.contains_literals = true;
                     self.elements
                         .push(UnionElement::BytesLiterals(FxOrderSet::from_iter([
                             literal,
@@ -452,10 +461,10 @@ impl<'db> UnionBuilder<'db> {
                 }
             }
             // And same for int literals as well.
-            Type::IntLiteral(literal) => {
+            Type::IntLiteral(literal) if self.contains_literals || self.elements.is_empty() => {
                 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) => {
@@ -473,8 +482,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();
@@ -487,6 +499,7 @@ impl<'db> UnionBuilder<'db> {
                 if let Some(found) = found {
                     found.insert(literal);
                 } else {
+                    self.contains_literals = true;
                     self.elements
                         .push(UnionElement::IntLiterals(FxOrderSet::from_iter([literal])));
                 }
@@ -549,19 +562,28 @@ impl<'db> UnionBuilder<'db> {
         // unpacking them.
         let should_simplify_full = !matches!(ty, Type::TypeAlias(_)) && !self.cycle_recovery;
 
-        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
+        let mut ty_negated: Option<Type> = None;
+
+        let mut i = 0;
+        let mut insertion_point: Option<usize> = None;
+
+        let mut remove_or_replace = |i: usize, elements: &mut Vec<UnionElement<'db>>| {
+            if insertion_point.is_none() {
+                insertion_point = Some(i);
+            } else {
+                elements.swap_remove(i);
+            }
         };
 
-        for (index, element) in self.elements.iter_mut().enumerate() {
+        while i < self.elements.len() {
+            let element = &mut self.elements[i];
+
             let element_type = match element.try_reduce(self.db, ty) {
                 ReduceResult::KeepIf(keep) => {
                     if !keep {
-                        to_remove.push(index);
+                        remove_or_replace(i, &mut self.elements);
                     }
+                    i += 1;
                     continue;
                 }
                 ReduceResult::Type(ty) => ty,
@@ -587,19 +609,24 @@ 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;
-                }
+            if element_type.is_typed_dict() && ty.is_typed_dict() {
+                i += 1;
                 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) {
+                    remove_or_replace(i, &mut self.elements);
+                    i += 1;
+                    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
@@ -613,13 +640,12 @@ impl<'db> UnionBuilder<'db> {
                     return;
                 }
             }
+
+            i += 1;
         }
-        if let Some((&first, rest)) = to_remove.split_first() {
-            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() {
-                self.elements.swap_remove(index);
-            }
+
+        if let Some(insertion_point) = insertion_point {
+            self.elements[insertion_point] = UnionElement::Type(ty);
         } else {
             self.elements.push(UnionElement::Type(ty));
         }