Avoid quadratic enum comparison narrowing

Author: charliermarsh

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md

@@ -186,6 +186,28 @@ def _(answer: IndependentEquality):
         reveal_type(answer)  # revealed: IndependentEquality
 ```
 
+Finite domains remain narrowable when the other operand also includes an identity singleton:
+
+```py
+from enum import Enum
+from typing import Literal
+
+class Finite(Enum):
+    FIRST = 1
+    SECOND = 2
+
+def _(value: Finite, other: Literal[Finite.FIRST] | None):
+    if value == other:
+        reveal_type(value)  # revealed: Literal[Finite.FIRST]
+    else:
+        reveal_type(value)  # revealed: Finite
+
+    if value != other:
+        reveal_type(value)  # revealed: Finite
+    else:
+        reveal_type(value)  # revealed: Literal[Finite.FIRST]
+```
+
 ## Equality between concrete runtime classes
 
 Types such as `bool`, `LiteralString`, and `TypedDict` correspond to specific runtime classes.

crates/ty_python_semantic/src/types/equality.rs

@@ -1,4 +1,5 @@
 use ruff_python_ast::name::Name;
+use rustc_hash::FxHashSet;
 
 use crate::{Db, place::PlaceAndQualifiers};
 
@@ -109,10 +110,39 @@ fn equality_result<'db>(
     let left = left.resolve_type_alias(db);
     let right = right.resolve_type_alias(db);
 
-    if let Some(alternatives) = equality_alternatives(db, left) {
+    if !finite_domain_expansion_is_bounded(db, left, right, ComparisonOperator::Equality) {
+        return ComparisonResult::Ambiguous;
+    }
+
+    let left_alternatives = finite_alternatives(db, left, ComparisonOperator::Equality);
+    if left == right
+        && let Some(alternatives) = left_alternatives.as_deref()
+    {
+        return if alternatives.len() == 1 {
+            ComparisonResult::AlwaysTrue
+        } else {
+            ComparisonResult::Ambiguous
+        };
+    }
+
+    let right_alternatives = finite_alternatives(db, right, ComparisonOperator::Equality);
+    if let (Some(left_alternatives), Some(right_alternatives)) =
+        (left_alternatives.as_deref(), right_alternatives.as_deref())
+    {
+        return evaluate_finite_domains(
+            db,
+            left_alternatives,
+            right_alternatives,
+            is_positive,
+            ComparisonOperator::Equality,
+            equality_result,
+        );
+    }
+
+    if let Some(alternatives) = left_alternatives {
         return evaluate_union_left(db, &alternatives, right, is_positive, equality_result);
     }
-    if let Some(alternatives) = equality_alternatives(db, right) {
+    if let Some(alternatives) = right_alternatives {
         return evaluate_union_right(db, left, &alternatives, is_positive, equality_result);
     }
 
@@ -378,10 +408,39 @@ fn inequality_result<'db>(
     let left = left.resolve_type_alias(db);
     let right = right.resolve_type_alias(db);
 
-    if let Some(alternatives) = inequality_alternatives(db, left) {
+    if !finite_domain_expansion_is_bounded(db, left, right, ComparisonOperator::Inequality) {
+        return ComparisonResult::Ambiguous;
+    }
+
+    let left_alternatives = finite_alternatives(db, left, ComparisonOperator::Inequality);
+    if left == right
+        && let Some(alternatives) = left_alternatives.as_deref()
+    {
+        return if alternatives.len() == 1 {
+            ComparisonResult::AlwaysFalse
+        } else {
+            ComparisonResult::Ambiguous
+        };
+    }
+
+    let right_alternatives = finite_alternatives(db, right, ComparisonOperator::Inequality);
+    if let (Some(left_alternatives), Some(right_alternatives)) =
+        (left_alternatives.as_deref(), right_alternatives.as_deref())
+    {
+        return evaluate_finite_domains(
+            db,
+            left_alternatives,
+            right_alternatives,
+            is_positive,
+            ComparisonOperator::Inequality,
+            inequality_result,
+        );
+    }
+
+    if let Some(alternatives) = left_alternatives {
         return evaluate_union_left(db, &alternatives, right, is_positive, inequality_result);
     }
-    if let Some(alternatives) = inequality_alternatives(db, right) {
+    if let Some(alternatives) = right_alternatives {
         return evaluate_union_right(db, left, &alternatives, is_positive, inequality_result);
     }
 
@@ -766,12 +825,167 @@ fn evaluate_intersection_left<'db>(
     }
 }
 
-fn equality_alternatives<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option<Vec<Type<'db>>> {
-    finite_alternatives(db, ty, ComparisonOperator::Equality)
+fn evaluate_finite_domains<'db>(
+    db: &'db dyn Db,
+    left: &[Type<'db>],
+    right: &[Type<'db>],
+    is_positive: bool,
+    operator: ComparisonOperator,
+    evaluate: fn(&'db dyn Db, Type<'db>, Type<'db>, bool) -> ComparisonResult<'db>,
+) -> ComparisonResult<'db> {
+    if left.is_empty() || right.is_empty() {
+        return ComparisonResult::Ambiguous;
+    }
+
+    let Some(other_keys) = right
+        .iter()
+        .map(|alternative| finite_comparison_key(db, *alternative, operator))
+        .collect::<Option<FxHashSet<_>>>()
+    else {
+        return ComparisonResult::Ambiguous;
+    };
+
+    evaluate_target_union(db, left, is_positive, |alternative| {
+        if let Some(key) = finite_comparison_key(db, alternative, operator) {
+            let equality = if !other_keys.contains(&key) {
+                ComparisonResult::AlwaysFalse
+            } else if other_keys.len() == 1 {
+                ComparisonResult::AlwaysTrue
+            } else {
+                ComparisonResult::Ambiguous
+            };
+            if operator == ComparisonOperator::Equality {
+                equality
+            } else {
+                equality.negate()
+            }
+        } else {
+            evaluate_against_results(
+                db,
+                alternative,
+                is_positive,
+                right
+                    .iter()
+                    .map(|other| evaluate(db, alternative, *other, is_positive)),
+            )
+        }
+    })
+}
+
+fn finite_comparison_key<'db>(
+    db: &'db dyn Db,
+    ty: Type<'db>,
+    operator: ComparisonOperator,
+) -> Option<Type<'db>> {
+    let literal = match ty {
+        Type::LiteralValue(literal) => literal.kind(),
+        Type::Intersection(intersection) => LiteralValueTypeKind::Enum(
+            intersection
+                .positive(db)
+                .iter()
+                .find_map(|element| element.as_enum_literal())?,
+        ),
+        _ if has_known_identity_comparison_semantics(db, ty, operator) => return Some(ty),
+        _ => return None,
+    };
+
+    match literal {
+        LiteralValueTypeKind::Bool(value) => Some(Type::int_literal(i64::from(value))),
+        LiteralValueTypeKind::Int(_)
+        | LiteralValueTypeKind::String(_)
+        | LiteralValueTypeKind::Bytes(_) => Some(ty),
+        LiteralValueTypeKind::LiteralString => None,
+        LiteralValueTypeKind::Enum(enum_literal) => {
+            match known_instance_semantics(db, enum_literal.enum_class_instance(db), operator)? {
+                KnownComparisonSemantics::Object => {
+                    let metadata = enum_metadata(db, enum_literal.enum_class(db))?;
+                    let name = metadata.resolve_member(enum_literal.name(db))?;
+                    Some(Type::enum_literal(EnumLiteralType::new(
+                        db,
+                        enum_literal.enum_class(db),
+                        name.clone(),
+                    )))
+                }
+                KnownComparisonSemantics::Int
+                | KnownComparisonSemantics::Str
+                | KnownComparisonSemantics::Bytes => {
+                    finite_comparison_key(db, enum_literal_value(db, enum_literal)?, operator)
+                }
+                KnownComparisonSemantics::Tuple | KnownComparisonSemantics::Dict => None,
+            }
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+enum FiniteComparisonDomain {
+    None,
+    Finite,
+    Mixed,
 }
 
-fn inequality_alternatives<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option<Vec<Type<'db>>> {
-    finite_alternatives(db, ty, ComparisonOperator::Inequality)
+fn finite_domain_expansion_is_bounded(
+    db: &dyn Db,
+    target: Type,
+    other: Type,
+    operator: ComparisonOperator,
+) -> bool {
+    !matches!(other, Type::Union(_))
+        || finite_comparison_domain(db, target, operator) == FiniteComparisonDomain::None
+        || finite_comparison_domain(db, other, operator) == FiniteComparisonDomain::Finite
+}
+
+fn finite_comparison_domain(
+    db: &dyn Db,
+    ty: Type,
+    operator: ComparisonOperator,
+) -> FiniteComparisonDomain {
+    match ty {
+        Type::Union(union) => {
+            let mut has_finite = false;
+            let mut has_open = false;
+            for element in union.elements(db) {
+                match finite_comparison_domain(db, *element, operator) {
+                    FiniteComparisonDomain::None => has_open = true,
+                    FiniteComparisonDomain::Finite => has_finite = true,
+                    FiniteComparisonDomain::Mixed => return FiniteComparisonDomain::Mixed,
+                }
+            }
+            match (has_finite, has_open) {
+                (true, true) => FiniteComparisonDomain::Mixed,
+                (true, false) => FiniteComparisonDomain::Finite,
+                (false, _) => FiniteComparisonDomain::None,
+            }
+        }
+        Type::EnumComplement(_) => {
+            if comparison_semantics(db, ty, operator).is_some() {
+                FiniteComparisonDomain::Finite
+            } else {
+                FiniteComparisonDomain::None
+            }
+        }
+        Type::Intersection(intersection) => {
+            if intersection.enum_complement(db).is_some()
+                && comparison_semantics(db, ty, operator).is_some()
+            {
+                FiniteComparisonDomain::Finite
+            } else {
+                FiniteComparisonDomain::None
+            }
+        }
+        _ if finite_comparison_key(db, ty, operator).is_some() => FiniteComparisonDomain::Finite,
+        Type::NominalInstance(instance) => {
+            if instance.has_known_class(db, KnownClass::Bool)
+                || (enum_metadata(db, instance.class_literal(db)).is_some()
+                    && comparison_semantics(db, ty, operator).is_some())
+            {
+                FiniteComparisonDomain::Finite
+            } else {
+                FiniteComparisonDomain::None
+            }
+        }
+        _ => FiniteComparisonDomain::None,
+    }
 }
 
 fn finite_alternatives<'db>(
@@ -780,6 +994,23 @@ fn finite_alternatives<'db>(
     operator: ComparisonOperator,
 ) -> Option<Vec<Type<'db>>> {
     match ty {
+        Type::Union(union) => {
+            let mut alternatives = Vec::new();
+            let mut expanded_finite_domain = false;
+            for element in union.elements(db) {
+                if let Some(element_alternatives) = finite_alternatives(db, *element, operator) {
+                    alternatives.extend(element_alternatives);
+                    expanded_finite_domain = true;
+                } else {
+                    alternatives.push(*element);
+                }
+            }
+            (expanded_finite_domain
+                || alternatives
+                    .iter()
+                    .all(|alternative| finite_comparison_key(db, *alternative, operator).is_some()))
+            .then_some(alternatives)
+        }
         Type::EnumComplement(complement) => comparison_semantics(db, ty, operator)
             .is_some()
             .then(|| complement.remaining_literal_types(db)),

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

@@ -113,6 +113,56 @@ def classify(message: Message) -> int:
     Ok(())
 }
 
+#[test]
+fn enum_comparison_narrowing_avoids_quadratic_expansion() -> anyhow::Result<()> {
+    let mut db = setup_db();
+    let mut members = String::new();
+    for index in 0..250 {
+        writeln!(&mut members, "    MEMBER_{index} = {index}")
+            .expect("writing to a String cannot fail");
+    }
+    let source = format!(
+        r#"
+from enum import Enum
+
+class Left(Enum):
+{members}
+class Right(Enum):
+{members}
+def consume(value: object) -> None: ...
+
+def compare_same(left: Left, right: Left) -> None:
+    if left == right:
+        consume(left)
+    if left != right:
+        consume(left)
+
+def compare_optional(left: Left, right: Left | None) -> None:
+    if left == right:
+        consume(left)
+    if left != right:
+        consume(left)
+
+def compare_different(left: Left, right: Right) -> None:
+    if left == right:
+        consume(left)
+    if left != right:
+        consume(left)
+"#
+    );
+    db.write_file("/src/a.py", source)?;
+
+    let start = Instant::now();
+    assert_file_diagnostics(&db, "/src/a.py", &[]);
+    assert!(
+        start.elapsed() < Duration::from_secs(10),
+        "enum comparison narrowing took {:?}",
+        start.elapsed()
+    );
+
+    Ok(())
+}
+
 #[test]
 fn not_literal_string() -> anyhow::Result<()> {
     let mut db = setup_db();