wip

Author: AlexWaygood

crates/red_knot_python_semantic/src/symbol.rs

@@ -549,10 +549,11 @@ fn symbol_from_bindings_impl<'db>(
         };
 
         if let Some(second) = types.next() {
-            Symbol::Type(
-                UnionType::from_elements(db, [first, second].into_iter().chain(types)),
-                boundness,
-            )
+            let ty = UnionBuilder::new(db)
+                .extend([first, second])
+                .extend(types)
+                .build();
+            Symbol::Type(ty, boundness)
         } else {
             Symbol::Type(first, boundness)
         }

crates/red_knot_python_semantic/src/types.rs

@@ -4172,16 +4172,15 @@ impl<'db> UnionType<'db> {
 
     /// Create a union from a list of elements
     /// (which may be eagerly simplified into a different variant of [`Type`] altogether).
-    pub fn from_elements<I, T>(db: &'db dyn Db, elements: I) -> Type<'db>
+    pub fn from_elements<I, E, T>(db: &'db dyn Db, elements: I) -> Type<'db>
     where
-        I: IntoIterator<Item = T>,
+        I: IntoIterator<IntoIter = E>,
+        E: ExactSizeIterator<Item = T>,
         T: Into<Type<'db>>,
     {
-        elements
-            .into_iter()
-            .fold(UnionBuilder::new(db), |builder, element| {
-                builder.add(element.into())
-            })
+        let elements = elements.into_iter();
+        UnionBuilder::with_capacity(db, elements.len())
+            .extend(elements)
             .build()
     }
 

crates/red_knot_python_semantic/src/types/builder.rs

@@ -43,6 +43,26 @@ impl<'db> UnionBuilder<'db> {
         }
     }
 
+    pub(crate) fn with_capacity(db: &'db dyn Db, capacity: usize) -> Self {
+        Self {
+            db,
+            elements: Vec::with_capacity(capacity),
+        }
+    }
+
+    pub(crate) fn extend<I, T>(mut self, types: I) -> Self
+    where
+        I: IntoIterator<Item = T>,
+        T: Into<Type<'db>>,
+    {
+        let types = types.into_iter();
+        self.elements.reserve(types.size_hint().0);
+        for ty in types {
+            self = self.add(ty.into());
+        }
+        self
+    }
+
     /// Collapse the union to a single type: `object`.
     fn collapse_to_object(mut self) -> Self {
         self.elements.clear();

crates/red_knot_python_semantic/src/types/call.rs

@@ -1,5 +1,5 @@
 use super::context::InferContext;
-use super::{Signature, Type};
+use super::{Signature, Type, UnionBuilder};
 use crate::types::UnionType;
 use crate::Db;
 
@@ -128,13 +128,13 @@ impl<'db> CallError<'db> {
             // combined with `Type::Unknown`
             CallError::Union {
                 errors, bindings, ..
-            } => Some(UnionType::from_elements(
-                db,
-                bindings
-                    .iter()
-                    .map(CallBinding::return_type)
-                    .chain(errors.iter().map(|err| err.fallback_return_type(db))),
-            )),
+            } => Some(
+                UnionBuilder::new(db)
+                    .extend(bindings.iter().map(CallBinding::return_type))
+                    .extend(errors.iter().map(|err| err.fallback_return_type(db)))
+                    .build(),
+            ),
+
             Self::PossiblyUnboundDunderCall { outcome, .. } => Some(outcome.return_type(db)),
             Self::BindingError { binding } => Some(binding.return_type()),
         }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -4079,7 +4079,9 @@ impl<'db> TypeInferenceBuilder<'db> {
             }
         });
 
-        UnionType::from_elements(db, elements)
+        UnionBuilder::with_capacity(db, n_values)
+            .extend(elements)
+            .build()
     }
 
     fn infer_compare_expression(&mut self, compare: &ast::ExprCompare) -> Type<'db> {

crates/red_knot_python_semantic/src/types/property_tests.rs

@@ -323,7 +323,11 @@ fn intersection<'db>(db: &'db TestDb, tys: impl IntoIterator<Item = Type<'db>>)
     builder.build()
 }
 
-fn union<'db>(db: &'db TestDb, tys: impl IntoIterator<Item = Type<'db>>) -> Type<'db> {
+fn union<'db, I, E>(db: &'db TestDb, tys: I) -> Type<'db>
+where
+    I: IntoIterator<IntoIter = E>,
+    E: ExactSizeIterator<Item = Type<'db>>,
+{
     UnionType::from_elements(db, tys)
 }