Speed up match checking after enum equality narrowing

mypy/meet.py

@@ -114,6 +114,66 @@ def meet_types(s: Type, t: Type) -> ProperType:
     return t.accept(TypeMeetVisitor(s))
 
 
+def _is_simple_union_item(t: ProperType) -> bool:
+    """Items that overlap another simple item iff they compare equal."""
+    return isinstance(t, (LiteralType, NoneType, UninhabitedType))
+
+
+def _union_items_may_narrow(declared: Type, narrowed: Type) -> bool:
+    """Keep the historical union-narrowing overlap predicate.
+
+    This special-casing is needed to support checking branches like this:
+
+        x: Union[float, complex]
+        if isinstance(x, int):
+            ...
+
+    And assignments like this:
+
+        x: float | None
+        y: int | None
+        x = y
+    """
+    return is_overlapping_types(declared, narrowed, ignore_promotions=True) or is_subtype(
+        narrowed, declared, ignore_promotions=False
+    )
+
+
+def _narrow_declared_union_items(
+    declared_items: list[Type], narrowed_items: list[Type]
+) -> ProperType:
+    """Intersect two unions, hashing simple items to avoid O(N*M) overlap checks.
+
+    After equality narrowing, enums explode into unions of literals. Intersecting
+    two such unions pairwise is quadratic and dominates match-statement checking
+    on large enums (see #21997). Literal/None/Never items can be intersected via
+    hashing; remaining items still use the original overlap predicate.
+    """
+    simple_narrowed: dict[ProperType, Type] = {}
+    complex_narrowed: list[Type] = []
+    for item in narrowed_items:
+        proper = get_proper_type(item)
+        if _is_simple_union_item(proper):
+            simple_narrowed[proper] = item
+        else:
+            complex_narrowed.append(item)
+
+    result: list[Type] = []
+    for declared in declared_items:
+        proper = get_proper_type(declared)
+        if _is_simple_union_item(proper):
+            if proper in simple_narrowed:
+                result.append(declared)
+            for narrowed in complex_narrowed:
+                if _union_items_may_narrow(declared, narrowed):
+                    result.append(narrow_declared_type(declared, narrowed))
+        else:
+            for narrowed in narrowed_items:
+                if _union_items_may_narrow(declared, narrowed):
+                    result.append(narrow_declared_type(declared, narrowed))
+    return make_simplified_union(result)
+
+
 def narrow_declared_type(declared: Type, narrowed: Type) -> Type:
     """Return the declared type narrowed down to another type."""
     # TODO: check infinite recursion for aliases here.
@@ -133,30 +193,12 @@ def narrow_declared_type(declared: Type, narrowed: Type) -> Type:
         return original_declared
     if isinstance(declared, UnionType):
         declared_items = declared.relevant_items()
+        narrowed_items: list[Type]
         if isinstance(narrowed, UnionType):
             narrowed_items = narrowed.relevant_items()
         else:
             narrowed_items = [narrowed]
-        return make_simplified_union(
-            [
-                narrow_declared_type(d, n)
-                for d in declared_items
-                for n in narrowed_items
-                # This (ugly) special-casing is needed to support checking
-                # branches like this:
-                # x: Union[float, complex]
-                # if isinstance(x, int):
-                #     ...
-                # And assignments like this:
-                # x: float | None
-                # y: int | None
-                # x = y
-                if (
-                    is_overlapping_types(d, n, ignore_promotions=True)
-                    or is_subtype(n, d, ignore_promotions=False)
-                )
-            ]
-        )
+        return _narrow_declared_union_items(declared_items, narrowed_items)
     if is_enum_overlapping_union(declared, narrowed):
         # Quick check before reaching `is_overlapping_types`. If it's enum/literal overlap,
         # avoid full expansion and make it faster.

mypy/test/testtypes.py

@@ -1359,6 +1359,7 @@ def test_type_type(self) -> None:
 
     def test_literal_type(self) -> None:
         a = self.fx.a
+        b = self.fx.b
         lit1 = self.fx.lit1
         lit2 = self.fx.lit2
         lit3 = self.fx.lit3
@@ -1376,6 +1377,30 @@ def test_literal_type(self) -> None:
         assert is_same_type(lit1, narrow_declared_type(lit1, a))
         assert is_same_type(lit2, narrow_declared_type(lit2, a))
 
+        # Union-union narrowing keeps overlapping literals and drops the rest.
+        assert is_same_type(
+            lit2, narrow_declared_type(UnionType([lit1, lit2]), UnionType([lit2, lit3]))
+        )
+        # Mixed literal/instance unions still use the overlap predicate.
+        assert is_same_type(
+            UnionType([lit1, b]), narrow_declared_type(UnionType([lit1, a]), UnionType([lit1, b]))
+        )
+
+    def test_narrow_declared_type_large_literal_unions(self) -> None:
+        # Equality narrowing explodes enums into large literal unions. Intersecting
+        # those unions used to be quadratic; this should stay effectively linear.
+        fx = self.fx
+        left = [LiteralType("v%d" % i, fx.str_type) for i in range(2000)]
+        right = [LiteralType("v%d" % i, fx.str_type) for i in range(500, 2500)]
+        result = get_proper_type(narrow_declared_type(UnionType(left), UnionType(right)))
+        assert isinstance(result, UnionType)
+        values = set()
+        for item in result.items:
+            proper = get_proper_type(item)
+            if isinstance(proper, LiteralType):
+                values.add(proper.value)
+        assert values == {"v%d" % i for i in range(500, 2000)}
+
     # FIX generic interfaces + ranges
 
     def assert_meet_uninhabited(self, s: Type, t: Type) -> None:

test-data/unit/check-python310.test

@@ -2310,6 +2310,25 @@ def g(m: Medal) -> int:
             return 2
 [builtins fixtures/enum.pyi]
 
+[case testMatchAfterEqualityNarrowingEnum]
+# flags: --strict-equality --warn-unreachable
+from enum import Enum
+
+class Medal(Enum):
+    gold = 1
+    silver = 2
+    bronze = 3
+
+def f(m: Medal) -> None:
+    if m == Medal.gold:
+        return
+    match m:
+        case Medal.silver:
+            reveal_type(m)  # N: Revealed type is "Literal[__main__.Medal.silver]"
+        case Medal.bronze:
+            reveal_type(m)  # N: Revealed type is "Literal[__main__.Medal.bronze]"
+[builtins fixtures/primitives.pyi]
+
 [case testMatchLiteralOrValuePattern]
 # flags: --strict-equality --warn-unreachable
 from typing import Literal