From d5089ccd4d28716a9fc38e4c26f71b98de2a02a2 Mon Sep 17 00:00:00 2001
From: Douglas Creager <dcreager@dcreager.net>
Date: Thu, 8 Jan 2026 09:05:23 -0500
Subject: [PATCH] estimate size of distributed upper bound

---
 crates/ty_python_semantic/src/types.rs        | 31 +++++++++++++++----
 .../src/types/constraints.rs                  | 31 ++++++++++---------
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs
index b26842222b..bc04bb2ed6 100644
--- a/crates/ty_python_semantic/src/types.rs
+++ b/crates/ty_python_semantic/src/types.rs
@@ -1359,12 +1359,6 @@ impl<'db> Type<'db> {
         self.as_union().expect("Expected a Type::Union variant")
     }
 
-    pub(crate) fn union_clause_count(self, db: &'db dyn Db) -> usize {
-        self.as_union()
-            .map(|union_type| union_type.elements(db).len())
-            .unwrap_or(1)
-    }
-
     /// Returns whether this is a "real" intersection type. (Negated types are represented by an
     /// intersection containing a single negative branch, which this method does _not_ consider a
     /// "real" intersection.)
@@ -1375,6 +1369,31 @@ impl<'db> Type<'db> {
         }
     }
 
+    /// Returns the number of union clauses in this type. If the type is not a union, returns 1.
+    pub(crate) fn union_size(self, db: &'db dyn Db) -> usize {
+        self.as_union()
+            .map(|union_type| union_type.elements(db).len())
+            .unwrap_or(1)
+    }
+
+    /// Returns the number of intersection clauses in this type. If the type is a union, this is
+    /// the maximum of the `intersection_size` of each union element. If the type is not a union
+    /// nor an intersection, returns 1.
+    pub(crate) fn intersection_size(self, db: &'db dyn Db) -> usize {
+        match self {
+            Type::Intersection(intersection) => {
+                intersection.positive(db).len() + intersection.negative(db).len()
+            }
+            Type::Union(union_type) => union_type
+                .elements(db)
+                .iter()
+                .map(|element| element.intersection_size(db))
+                .max()
+                .unwrap_or(1),
+            _ => 1,
+        }
+    }
+
     pub(crate) const fn as_function_literal(self) -> Option<FunctionType<'db>> {
         match self {
             Type::FunctionLiteral(function_type) => Some(function_type),
diff --git a/crates/ty_python_semantic/src/types/constraints.rs b/crates/ty_python_semantic/src/types/constraints.rs
index 6148cb59af..6162f9f887 100644
--- a/crates/ty_python_semantic/src/types/constraints.rs
+++ b/crates/ty_python_semantic/src/types/constraints.rs
@@ -758,9 +758,25 @@ impl<'db> ConstrainedTypeVar<'db> {
 
     /// Returns the intersection of two range constraints, or `None` if the intersection is empty.
     fn intersect(self, db: &'db dyn Db, other: Self) -> IntersectionResult<'db> {
+        // TODO: For now, we treat some upper bounds as unsimplifiable if they become "too big".
+        // When intersecting constraints, the upper bounds are also intersected together. If the
+        // lhs and rhs upper bounds are unions of intersections (e.g. `(a & b) | (c & d)`), then
+        // intersecting them together will require distributing across every pair of union
+        // elements. That can quickly balloon in size. We are looking at a better representation
+        // that would let us model this case more directly, but for now, we punt.
+        let self_upper = self.upper(db);
+        let other_upper = other.upper(db);
+        let estimated_upper_bound_size = self_upper.union_size(db)
+            * other_upper.union_size(db)
+            * (self_upper.intersection_size(db) + other_upper.intersection_size(db));
+        const MAX_UPPER_BOUND_SIZE: usize = 4;
+        if estimated_upper_bound_size >= MAX_UPPER_BOUND_SIZE {
+            return IntersectionResult::CannotSimplify;
+        }
+
         // (s₁ ≤ α ≤ t₁) ∧ (s₂ ≤ α ≤ t₂) = (s₁ ∪ s₂) ≤ α ≤ (t₁ ∩ t₂))
         let lower = UnionType::from_elements(db, [self.lower(db), other.lower(db)]);
-        let upper = IntersectionType::from_elements(db, [self.upper(db), other.upper(db)]);
+        let upper = IntersectionType::from_elements(db, [self_upper, other_upper]);
 
         // If `lower ≰ upper`, then the intersection is empty, since there is no type that is both
         // greater than `lower`, and less than `upper`.
@@ -774,19 +790,6 @@ impl<'db> ConstrainedTypeVar<'db> {
             return IntersectionResult::CannotSimplify;
         }
 
-        // TODO: For now, we also treat upper bound unions as unsimplifiable if they become too
-        // big. Upper bounds are intersected together, and the intersections of large unions can
-        // become quite large indeed. We are looking at a better representation that would let us
-        // model them directly, but for now, we punt. Instead of hard-coding a specific size
-        // threshold, we skip any upper bounds that are "larger" than either of constraints being
-        // intersected.
-        let upper_size = upper.union_clause_count(db);
-        if upper_size > self.upper(db).union_clause_count(db)
-            || upper_size > other.upper(db).union_clause_count(db)
-        {
-            return IntersectionResult::CannotSimplify;
-        }
-
         IntersectionResult::Simplified(Self::new(db, self.typevar(db), lower, upper))
     }