typeops: extend make_simplified_union fast path to enums

In PR python#9192 a fast path was created to address the slowness reported
in issue python#9169 wherein large Union or literal types would dramatically
slow down typechecking.

It is desirable to extend this fast path to cover Enum types, as these
can also leverage the O(n) set-based fast path instead of the O(n**2)
fallback.

This is seen to bring down the typechecking of a single fairly simple
chain of `if` statements operating on a large enum (~3k members) from
~40min to 12s in real-world code! Note that the timing is taken from
a pure-python run of mypy, as opposed to a compiled version.

mypy/typeops.py

@@ -5,7 +5,7 @@
       since these may assume that MROs are ready.
 """
 
-from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar
+from typing import cast, Optional, List, Sequence, Set, Iterable, TypeVar, Tuple
 from typing_extensions import Type as TypingType
 import sys
 
@@ -344,35 +344,41 @@ def make_simplified_union(items: Sequence[Type],
     from mypy.subtypes import is_proper_subtype
 
     removed = set()  # type: Set[int]
-
-    # Avoid slow nested for loop for Union of Literal of strings (issue #9169)
-    if all((isinstance(item, LiteralType) and
-            item.fallback.type.fullname == 'builtins.str')
-           for item in items):
-        seen = set()    # type: Set[str]
-        for index, item in enumerate(items):
+    seen_fast = set()  # type: Set[Tuple[str, str]]
+
+    # NB: having a separate fast path for Union of Literal and slow path for other things
+    # would arguably be cleaner, however it breaks down when simplifying the Union of two
+    # different enum types as try_expanding_enum_to_union works recursively and will
+    # trigger intermediate simplifications that would render the fast path useless
+    # Besides, the blended loop allows us to gracefully handle unions involving none/optional
+    # without any extra logic.
+    for i, item in enumerate(items):
+        if i in removed:
+            continue
+        # Avoid slow nested for loop for Union of Literal of strings/enums (issue #9169)
+        if isinstance(item, LiteralType) and (
+            item.fallback.type.is_enum or item.fallback.type.fullname == 'builtins.str'
+        ):
             assert isinstance(item, LiteralType)
             assert isinstance(item.value, str)
-            if item.value in seen:
-                removed.add(index)
-            seen.add(item.value)
-
-    else:
-        for i, ti in enumerate(items):
-            if i in removed: continue
-            # Keep track of the truishness info for deleted subtypes which can be relevant
-            cbt = cbf = False
-            for j, tj in enumerate(items):
-                if i != j and is_proper_subtype(tj, ti, keep_erased_types=keep_erased):
-                    # We found a redundant item in the union.
-                    removed.add(j)
-                    cbt = cbt or tj.can_be_true
-                    cbf = cbf or tj.can_be_false
-            # if deleted subtypes had more general truthiness, use that
-            if not ti.can_be_true and cbt:
-                items[i] = true_or_false(ti)
-            elif not ti.can_be_false and cbf:
-                items[i] = true_or_false(ti)
+            k = (item.value, item.fallback.type.fullname)
+            if k in seen_fast:
+                removed.add(i)
+            seen_fast.add(k)
+            continue
+        # Keep track of the truishness info for deleted subtypes which can be relevant
+        cbt = cbf = False
+        for j, tj in enumerate(items[i+1:], start=i+1):
+            if not isinstance(tj, LiteralType) and is_proper_subtype(tj, item, keep_erased_types=keep_erased):
+                # We found a redundant item in the union.
+                removed.add(j)
+                cbt = cbt or tj.can_be_true
+                cbf = cbf or tj.can_be_false
+        # if deleted subtypes had more general truthiness, use that
+        if not item.can_be_true and cbt:
+            items[i] = true_or_false(item)
+        elif not item.can_be_false and cbf:
+            items[i] = true_or_false(item)
 
     simplified_set = [items[i] for i in range(len(items)) if i not in removed]
     return UnionType.make_union(simplified_set, line, column)