[ty] Lazily build TypeVar accumulations (#24782)

## Summary

Given a generic specialization, we were rebuilding the constraints after
every argument, rather than all-at-once.

E.g., for:

```python
def combine[T](a: T, b: T, c: T, d: T) -> T:
    return a

combine(("name", 1), ("id", 2), ("flag", True), ("size", 4))
```

Each argument constrains the same type variable `T`:

```python
T = tuple[Literal["name"], Literal[1]]
T = tuple[Literal["id"], Literal[2]]
T = tuple[Literal["flag"], Literal[True]]
T = tuple[Literal["size"], Literal[4]]
```

On main, we then compute (roughly):
```
T = A
T = union(A, B)
T = union(union(A, B), C)
T = union(union(A, B, C), D)
```

Now, we create a builder and construct at the end. This has a
significant impact on functions with many arguments, but also reduces
memory on real-world projects, which is great.

crates/ty_python_semantic/src/types.rs

@@ -34,7 +34,9 @@ pub(crate) use self::iteration::extract_fixed_length_iterable_element_types;
 pub use self::known_instance::KnownInstanceType;
 pub(crate) use self::relation_error::{ErrorContext, ErrorContextTree, ParameterDescription};
 use self::set_theoretic::KnownUnion;
-pub(crate) use self::set_theoretic::builder::{IntersectionBuilder, UnionBuilder};
+pub(crate) use self::set_theoretic::builder::{
+    IntersectionBuilder, UnionAccumulator, UnionBuilder,
+};
 pub use self::set_theoretic::{
     IntersectionType, NegativeIntersectionElements, NegativeIntersectionElementsIterator, UnionType,
 };

crates/ty_python_semantic/src/types/call/bind.rs

@@ -55,7 +55,8 @@ use crate::types::{
     DataclassFlags, DataclassParams, GenericAlias, InternedConstraintSet, IntersectionType,
     KnownBoundMethodType, KnownClass, KnownInstanceType, LiteralValueTypeKind, NominalInstanceType,
     PropertyInstanceType, SpecialFormType, TypeAliasType, TypeContext, TypeVarBoundOrConstraints,
-    TypeVarVariance, UnionBuilder, UnionType, WrapperDescriptorKind, enums, list_members,
+    TypeVarVariance, UnionAccumulator, UnionBuilder, UnionType, WrapperDescriptorKind, enums,
+    list_members,
 };
 use crate::{DisplaySettings, FxOrderSet, Program};
 use ruff_db::diagnostic::{Annotation, Diagnostic, Span, SubDiagnostic, SubDiagnosticSeverity};
@@ -4347,7 +4348,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
                     return None;
                 };
 
-                let mut preferred: FxHashMap<BoundTypeVarIdentity<'db>, Type<'db>> =
+                let mut preferred: FxHashMap<BoundTypeVarIdentity<'db>, UnionAccumulator<'db>> =
                     FxHashMap::default();
 
                 for solution in &solutions {
@@ -4391,14 +4392,16 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
 
                         preferred
                             .entry(identity)
-                            .and_modify(|existing| {
-                                *existing =
-                                    UnionType::from_two_elements(self.db, *existing, inferred_ty);
-                            })
-                            .or_insert(inferred_ty);
+                            .and_modify(|existing| existing.add(self.db, inferred_ty))
+                            .or_insert_with(|| UnionAccumulator::new(inferred_ty));
                     }
                 }
 
+                let preferred: FxHashMap<BoundTypeVarIdentity<'db>, Type<'db>> = preferred
+                    .into_iter()
+                    .map(|(identity, accumulator)| (identity, accumulator.into_type(self.db)))
+                    .collect();
+
                 // Add preferred types to the builder so they serve as the base mapping
                 // when argument inference adds more types.
                 for solution in &solutions {

crates/ty_python_semantic/src/types/generics.rs

@@ -29,7 +29,7 @@ use crate::types::{
     ApplyTypeMappingVisitor, BindingContext, BoundTypeVarInstance, CallableType, CallableTypes,
     ClassLiteral, FindLegacyTypeVarsVisitor, IntersectionType, KnownClass, KnownInstanceType,
     MaterializationKind, Type, TypeAliasType, TypeContext, TypeMapping, TypeVarBoundOrConstraints,
-    TypeVarKind, TypeVarVariance, UnionType, binding_type, declaration_type,
+    TypeVarKind, TypeVarVariance, UnionAccumulator, UnionType, binding_type, declaration_type,
     infer_definition_types,
 };
 use crate::{Db, FxIndexMap, FxOrderMap, FxOrderSet};
@@ -1698,7 +1698,7 @@ pub(crate) struct SpecializationBuilder<'db, 'c> {
     db: &'db dyn Db,
     constraints: &'c ConstraintSetBuilder<'db>,
     inferable: InferableTypeVars<'db>,
-    types: FxHashMap<BoundTypeVarIdentity<'db>, Type<'db>>,
+    types: FxHashMap<BoundTypeVarIdentity<'db>, UnionAccumulator<'db>>,
 }
 
 /// An assignment from a bound type variable to a given type, along with the variance of the outermost
@@ -1737,12 +1737,14 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> {
             Option<(Type<'db>, Type<'db>)>,
         ) -> Option<Type<'db>>,
     ) -> Specialization<'db> {
+        let db = self.db;
         let types = generic_context
-            .variables_inner(self.db)
+            .variables_inner(db)
             .iter()
             .map(|(identity, variable)| {
-                match self.types.get(identity).copied() {
+                match self.types.get_mut(identity) {
                     Some(mapped_ty) => {
+                        let mapped_ty = mapped_ty.get_or_build(db);
                         // The typevar was inferred — present both bounds as the inferred type.
                         let chosen = choose(*variable, Some((mapped_ty, mapped_ty)));
                         Some(chosen.unwrap_or(mapped_ty))
@@ -1751,7 +1753,7 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> {
                 }
             });
 
-        generic_context.specialize_recursive(self.db, types)
+        generic_context.specialize_recursive(db, types)
     }
 
     /// Insert a type mapping for a bound typevar.
@@ -1776,10 +1778,10 @@ impl<'db, 'c> SpecializationBuilder<'db, 'c> {
                     return;
                 }
 
-                *entry.get_mut() = UnionType::from_two_elements(self.db, *entry.get(), ty);
+                entry.get_mut().add(self.db, ty);
             }
             Entry::Vacant(entry) => {
-                entry.insert(ty);
+                entry.insert(UnionAccumulator::new(ty));
             }
         }
     }

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -93,8 +93,9 @@ use crate::types::{
     KnownClass, KnownInstanceType, KnownUnion, LiteralValueTypeKind, MemberLookupPolicy,
     ParamSpecAttrKind, Parameter, ParameterForm, Parameters, Signature, SpecialFormType,
     SubclassOfType, Type, TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers,
-    TypeVarBoundOrConstraints, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder,
-    UnionType, binding_type, infer_complete_scope_types, infer_scope_types, todo_type,
+    TypeVarBoundOrConstraints, TypeVarKind, TypeVarVariance, TypedDictType, UnionAccumulator,
+    UnionBuilder, UnionType, binding_type, infer_complete_scope_types, infer_scope_types,
+    todo_type,
 };
 use crate::{AnalysisSettings, Db, FxIndexSet, Program};
 use ty_python_core::ast_ids::ScopedUseId;
@@ -6014,8 +6015,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         // type context is a covariant superclass of the collection (e.g., `Sequence[Any]` as type
         // context for `list[T]`).
         let (elt_tcx_constraints, elt_tcx_variance) = {
-            let mut elt_tcx_constraints: FxHashMap<BoundTypeVarIdentity<'_>, Type<'db>> =
-                FxHashMap::default();
+            let mut elt_tcx_constraints: FxHashMap<
+                BoundTypeVarIdentity<'db>,
+                UnionAccumulator<'db>,
+            > = FxHashMap::default();
             let mut elt_tcx_variance: FxHashMap<BoundTypeVarIdentity<'_>, TypeVarVariance> =
                 FxHashMap::default();
 
@@ -6081,14 +6084,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 let identity = binding.bound_typevar.identity(db);
                                 elt_tcx_constraints
                                     .entry(identity)
-                                    .and_modify(|existing| {
-                                        *existing = UnionType::from_two_elements(
-                                            db,
-                                            *existing,
-                                            inferred_ty,
-                                        );
-                                    })
-                                    .or_insert(inferred_ty);
+                                    .and_modify(|existing| existing.add(db, inferred_ty))
+                                    .or_insert_with(|| UnionAccumulator::new(inferred_ty));
                             }
                         }
 
@@ -6101,6 +6098,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 }
             }
 
+            let db = self.db();
+            let elt_tcx_constraints: FxHashMap<BoundTypeVarIdentity<'db>, Type<'db>> =
+                elt_tcx_constraints
+                    .into_iter()
+                    .map(|(identity, accumulator)| (identity, accumulator.into_type(db)))
+                    .collect();
+
             (elt_tcx_constraints, elt_tcx_variance)
         };
 

crates/ty_python_semantic/src/types/set_theoretic/builder.rs

@@ -350,6 +350,62 @@ pub(crate) struct UnionBuilder<'db> {
     recursively_defined: RecursivelyDefined,
 }
 
+/// Accumulates types into a union.
+///
+/// Most real-world type variables only accumulate one or two constraints. We keep those cases as
+/// plain `Type`s and only allocate a `UnionBuilder` once we know the accumulation is larger.
+pub(crate) enum UnionAccumulator<'db> {
+    One(Type<'db>),
+    Two(Type<'db>, Type<'db>),
+    Deferred(UnionBuilder<'db>),
+}
+
+impl<'db> UnionAccumulator<'db> {
+    pub(crate) fn new(ty: Type<'db>) -> Self {
+        UnionAccumulator::One(ty)
+    }
+
+    pub(crate) fn add(&mut self, db: &'db dyn Db, ty: Type<'db>) {
+        match self {
+            UnionAccumulator::One(existing) => {
+                *self = UnionAccumulator::Two(*existing, ty);
+            }
+            UnionAccumulator::Two(first, second) => {
+                let mut builder = UnionBuilder::new(db);
+                builder.add_in_place(*first);
+                builder.add_in_place(*second);
+                builder.add_in_place(ty);
+                *self = UnionAccumulator::Deferred(builder);
+            }
+            UnionAccumulator::Deferred(builder) => builder.add_in_place(ty),
+        }
+    }
+
+    pub(crate) fn get_or_build(&mut self, db: &'db dyn Db) -> Type<'db> {
+        match self {
+            UnionAccumulator::One(ty) => *ty,
+            UnionAccumulator::Two(first, second) => {
+                let ty = UnionType::from_two_elements(db, *first, *second);
+                *self = UnionAccumulator::One(ty);
+                ty
+            }
+            UnionAccumulator::Deferred(_) => {
+                let ty = std::mem::replace(self, UnionAccumulator::new(Type::Never)).into_type(db);
+                *self = UnionAccumulator::new(ty);
+                ty
+            }
+        }
+    }
+
+    pub(crate) fn into_type(self, db: &'db dyn Db) -> Type<'db> {
+        match self {
+            UnionAccumulator::One(ty) => ty,
+            UnionAccumulator::Two(first, second) => UnionType::from_two_elements(db, first, second),
+            UnionAccumulator::Deferred(builder) => builder.build(),
+        }
+    }
+}
+
 impl<'db> UnionBuilder<'db> {
     pub(crate) fn new(db: &'db dyn Db) -> Self {
         Self {