Skip to content

Conversation

@MichaReiser
Copy link
Member

Summary

UnionBuilder::from_elements is called very frequently. Let's see if short-circuiting in the case where elements is zero or one elements
long helps improve perf.

Test Plan

This should not change behavior.

@MichaReiser MichaReiser added the ty Multi-file analysis & type inference label Feb 17, 2025
@MichaReiser MichaReiser marked this pull request as draft February 17, 2025 17:38
@AlexWaygood
Copy link
Member

This looks great to me, and codspeed's reporting a 1% speedup. Another optimisation idea could be to pre-size the UnionBuilder, e.g.

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 4601543d4..6df55b467 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -4695,7 +4695,13 @@ impl<'db> UnionType<'db> {
             return first.into();
         };
 
-        let mut builder = UnionBuilder::new(db).add(first.into()).add(second.into());
+        let iterator_length = elements.size_hint().1.and_then(|size| size.checked_add(2));
+
+        let mut builder = iterator_length
+            .map(|num_elements| UnionBuilder::with_capacity(db, num_elements))
+            .unwrap_or_else(|| UnionBuilder::new(db))
+            .add(first.into())
+            .add(second.into());
 
         for element in elements {
             builder = builder.add(element.into());
diff --git a/crates/red_knot_python_semantic/src/types/builder.rs b/crates/red_knot_python_semantic/src/types/builder.rs
index 45be6ac48..c5cdf294d 100644
--- a/crates/red_knot_python_semantic/src/types/builder.rs
+++ b/crates/red_knot_python_semantic/src/types/builder.rs
@@ -43,6 +43,13 @@ impl<'db> UnionBuilder<'db> {
         }
     }
 
+    pub(crate) fn with_capacity(db: &'db dyn Db, capacity: usize) -> Self {
+        Self {
+            db,
+            elements: Vec::with_capacity(capacity),
+        }
+    }
+
     /// Collapse the union to a single type: `object`.
     fn collapse_to_object(mut self) -> Self {
         self.elements.clear();

@AlexWaygood AlexWaygood changed the title [red-knot] Avoid creating UnionBuilders unnecessary [red-knot] Avoid creating UnionBuilders unnecessarily Feb 18, 2025
@MichaReiser MichaReiser force-pushed the micha/union-builder-from-elements branch 2 times, most recently from 5ad6a67 to fe308d6 Compare February 18, 2025 13:17
@MichaReiser MichaReiser force-pushed the micha/union-builder-from-elements branch from fe308d6 to 6ce0a3c Compare February 18, 2025 13:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ty Multi-file analysis & type inference

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants