Skip to content
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.

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

[red-knot] Avoid creating UnionBuilders unnecessarily #16218

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from

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 red-knot 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
red-knot Multi-file analysis & type inference
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants